Caslow commited on
Commit
c49fe71
·
1 Parent(s): 6e308de

modify from_pretrained param

Browse files
Files changed (1) hide show
  1. inference.py +1 -3
inference.py CHANGED
@@ -26,9 +26,7 @@ def load_model(
26
  tokenizer = AutoTokenizer.from_pretrained(model_name)
27
 
28
  model = AutoModelForCausalLM.from_pretrained(
29
- model_name=model_name,
30
- torch_dtype=dtype,
31
- device_map="auto"
32
  )
33
 
34
  model.eval() # Set model to evaluation mode
 
26
  tokenizer = AutoTokenizer.from_pretrained(model_name)
27
 
28
  model = AutoModelForCausalLM.from_pretrained(
29
+ pretrained_model_name_or_path=model_name
 
 
30
  )
31
 
32
  model.eval() # Set model to evaluation mode