mew77 commited on
Commit
a8b4dfc
·
verified ·
1 Parent(s): 5ef8f13

Update hf_model.py

Browse files
Files changed (1) hide show
  1. hf_model.py +1 -1
hf_model.py CHANGED
@@ -13,7 +13,7 @@ class HFModel:
13
 
14
  try:
15
  log_info(f"=== Loading Model: {self.friendly_name} ===")
16
- self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()
17
  self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
18
  log_info(f"=== Model Loaded Successfully: {self.friendly_name} ===")
19
  except Exception as e:
 
13
 
14
  try:
15
  log_info(f"=== Loading Model: {self.friendly_name} ===")
16
+ self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16) #.cuda()
17
  self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
18
  log_info(f"=== Model Loaded Successfully: {self.friendly_name} ===")
19
  except Exception as e: