GMFTBY commited on
Commit
89690cf
β€’
1 Parent(s): fd75694

Update model/openllama.py

Browse files
Files changed (1) hide show
  1. model/openllama.py +1 -1
model/openllama.py CHANGED
@@ -116,7 +116,7 @@ class OpenLLAMAPEFTModel(nn.Module):
116
  )
117
 
118
  self.max_tgt_len = max_tgt_len
119
- self.device = torch.cuda.current_device()
120
 
121
  def encode_video(self, video_paths):
122
  inputs = {ModalityType.VISION: data.load_and_transform_video_data(video_paths, self.device)}
 
116
  )
117
 
118
  self.max_tgt_len = max_tgt_len
119
+ self.device = torch.cuda.current_device() if torch.cuda.is_available() else torch.device('cpu')
120
 
121
  def encode_video(self, video_paths):
122
  inputs = {ModalityType.VISION: data.load_and_transform_video_data(video_paths, self.device)}