Update README.md
Browse files
README.md
CHANGED
|
@@ -45,7 +45,7 @@ model = AutoModelForCausalLM.from_pretrained(
|
|
| 45 |
torch_dtype=torch.float16,
|
| 46 |
attn_implementation="flash_attention_2",
|
| 47 |
device_map="auto",
|
| 48 |
-
)
|
| 49 |
|
| 50 |
model = PeftModel.from_pretrained(model, model_id="henryen/OriGen")
|
| 51 |
model.eval()
|
|
|
|
| 45 |
torch_dtype=torch.float16,
|
| 46 |
attn_implementation="flash_attention_2",
|
| 47 |
device_map="auto",
|
| 48 |
+
).to("cuda")
|
| 49 |
|
| 50 |
model = PeftModel.from_pretrained(model, model_id="henryen/OriGen")
|
| 51 |
model.eval()
|