Update model/openllama.py

model/openllama.py

@@ -165,6 +165,8 @@ class OpenLLAMAPEFTModel(nn.Module):
         max_tgt_len = args['max_tgt_len']
         stage = args['stage']
 
+        self.device = torch.cuda.current_device()
+
         print (f'Initializing visual encoder from {imagebind_ckpt_path} ...')
 
         self.visual_encoder, self.visual_hidden_size = imagebind_model.imagebind_huge(args)
@@ -173,12 +175,12 @@ class OpenLLAMAPEFTModel(nn.Module):
 
         self.iter = 0
 
-        self.image_decoder = LinearLayer(1280, 1024, 4)
+        self.image_decoder = LinearLayer(1280, 1024, 4).to(self.device)
 
-        self.prompt_learner = PromptLearner(1, 4096)
+        self.prompt_learner = PromptLearner(1, 4096).to(self.device)
 
-        self.loss_focal = FocalLoss()
-        self.loss_dice = BinaryDiceLoss()
+        self.loss_focal = FocalLoss().to(self.device)
+        self.loss_dice = BinaryDiceLoss().to(self.device)
 
 
         # free vision encoder
@@ -213,7 +215,7 @@ class OpenLLAMAPEFTModel(nn.Module):
         )
 
         self.max_tgt_len = max_tgt_len
-        self.device = torch.cuda.current_device()
+        
 
 
     def rot90_img(self,x,k):