Spaces:
Sleeping
Sleeping
modify from_pretrained param
Browse files- 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 |
-
|
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
|