Update model/openllama.py

model/openllama.py

@@ -172,15 +172,16 @@ class OpenLLAMAPEFTModel(nn.Module):
         self.visual_encoder, self.visual_hidden_size = imagebind_model.imagebind_huge(args)
         imagebind_ckpt = torch.load(imagebind_ckpt_path, map_location=torch.device('cpu'))
         self.visual_encoder.load_state_dict(imagebind_ckpt, strict=True)
+        self.visual_encoder.to(self.device)
 
         self.iter = 0
 
-        self.image_decoder = LinearLayer(1280, 1024, 4).to(self.device)
+        self.image_decoder = LinearLayer(1280, 1024, 4)
 
-        self.prompt_learner = PromptLearner(1, 4096).to(self.device)
+        self.prompt_learner = PromptLearner(1, 4096)
 
-        self.loss_focal = FocalLoss().to(self.device)
-        self.loss_dice = BinaryDiceLoss().to(self.device)
+        self.loss_focal = FocalLoss()
+        self.loss_dice = BinaryDiceLoss()
 
 
         # free vision encoder