Update model/openllama.py

model/openllama.py

@@ -225,14 +225,14 @@ class OpenLLAMAPEFTModel(nn.Module):
         # self.llama_model.load_state_dict(delta_ckpt, strict=False)
         self.llama_model.print_trainable_parameters()
 
-        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16)
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16).to(self.device)
         self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
         self.llama_tokenizer.padding_side = "right"
         print ('Language decoder initialized.')
 
         self.llama_proj = nn.Linear(
             self.visual_hidden_size, self.llama_model.config.hidden_size
-        ).to(self.device)
+        ).to(torch.float16).to(self.device)
 
         self.max_tgt_len = max_tgt_len