Update README.md (#11)
Browse files- Update README.md (3f0f2cc0f25fa95e6b33c68c0df71bb7199f3445)
README.md
CHANGED
|
@@ -81,8 +81,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
| 81 |
model = AutoModelForCausalLM.from_pretrained(
|
| 82 |
model_name,
|
| 83 |
device_map="auto",
|
| 84 |
-
torch_dtype=torch.bfloat16
|
| 85 |
-
tp_plan="auto",
|
| 86 |
)
|
| 87 |
|
| 88 |
# prepare the model input
|
|
|
|
| 81 |
model = AutoModelForCausalLM.from_pretrained(
|
| 82 |
model_name,
|
| 83 |
device_map="auto",
|
| 84 |
+
torch_dtype=torch.bfloat16
|
|
|
|
| 85 |
)
|
| 86 |
|
| 87 |
# prepare the model input
|