Update README.md
Browse files
README.md
CHANGED
|
@@ -181,7 +181,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
| 181 |
model = AutoModelForCausalLM.from_pretrained(
|
| 182 |
model_name,
|
| 183 |
torch_dtype="auto",
|
| 184 |
-
device_map="
|
| 185 |
)
|
| 186 |
|
| 187 |
torch.cuda.reset_peak_memory_stats()
|
|
|
|
| 181 |
model = AutoModelForCausalLM.from_pretrained(
|
| 182 |
model_name,
|
| 183 |
torch_dtype="auto",
|
| 184 |
+
device_map="cuda:0"
|
| 185 |
)
|
| 186 |
|
| 187 |
torch.cuda.reset_peak_memory_stats()
|