Update model/openllama.py

model/openllama.py

@@ -170,15 +170,16 @@ class OpenLLAMAPEFTModel(nn.Module):
         print (f'Initializing visual encoder from {imagebind_ckpt_path} ...')
 
         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)
+        imagebind_ckpt = torch.load(imagebind_ckpt_path, map_location=self.device)
+        self.visual_encoder.load_state_dict(imagebind_ckpt, strict=True)
+        
 
         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()
@@ -202,11 +203,11 @@ class OpenLLAMAPEFTModel(nn.Module):
             target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']
         )
 
-        self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, device_map='auto', offload_folder="offload", offload_state_dict = True)
+        self.llama_model = LlamaForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', offload_folder="offload", offload_state_dict = True)
         self.llama_model = get_peft_model(self.llama_model, peft_config)
         self.llama_model.print_trainable_parameters()
 
-        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.bfloat16, device_map='auto', offload_folder="offload", offload_state_dict = True)
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16, device_map='auto', offload_folder="offload", offload_state_dict = True)
         self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
         self.llama_tokenizer.padding_side = "right"
         print ('Language decoder initialized.')