FantasticGNU commited on
Commit
9f4c400
1 Parent(s): 3c99987

Update model/openllama.py

Browse files
Files changed (1) hide show
  1. model/openllama.py +2 -2
model/openllama.py CHANGED
@@ -173,7 +173,7 @@ class OpenLLAMAPEFTModel(nn.Module):
173
 
174
  self.visual_encoder, self.visual_hidden_size = imagebind_model.imagebind_huge(args)
175
  self.visual_encoder.to(self.device)
176
- imagebind_ckpt = torch.load(imagebind_ckpt_path, map_location=torch.device('cuda'))
177
  self.visual_encoder.load_state_dict(imagebind_ckpt, strict=True)
178
 
179
 
@@ -215,7 +215,7 @@ class OpenLLAMAPEFTModel(nn.Module):
215
  # self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map=device_map, offload_folder="offload", offload_state_dict = True)
216
  # self.llama_model.to(torch.float16)
217
  # try:
218
- self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map=device_map, offload_folder="offload", offload_state_dict = True)
219
  # except:
220
  # pass
221
  # finally:
 
173
 
174
  self.visual_encoder, self.visual_hidden_size = imagebind_model.imagebind_huge(args)
175
  self.visual_encoder.to(self.device)
176
+ imagebind_ckpt = torch.load(imagebind_ckpt_path, map_location=torch.device('cpu'))
177
  self.visual_encoder.load_state_dict(imagebind_ckpt, strict=True)
178
 
179
 
 
215
  # self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map=device_map, offload_folder="offload", offload_state_dict = True)
216
  # self.llama_model.to(torch.float16)
217
  # try:
218
+ self.llama_model = AutoModelForCausalLM.from_pretrained('../vicuna_ckpt', torch_dtype=torch.float16, device_map=device_map, offload_folder="offload", offload_state_dict = True)
219
  # except:
220
  # pass
221
  # finally: