Update README.md
Browse files
README.md
CHANGED
@@ -59,8 +59,8 @@ if model_path.endswith("4bit"):
|
|
59 |
model = AutoModelForCausalLM.from_pretrained(
|
60 |
model_path,
|
61 |
load_in_4bit=True,
|
62 |
-
|
63 |
-
|
64 |
)
|
65 |
else:
|
66 |
model = AutoModelForCausalLM.from_pretrained(model_path).half().cuda()
|
|
|
59 |
model = AutoModelForCausalLM.from_pretrained(
|
60 |
model_path,
|
61 |
load_in_4bit=True,
|
62 |
+
torch_dtype=torch.float16,
|
63 |
+
device_map='auto'
|
64 |
)
|
65 |
else:
|
66 |
model = AutoModelForCausalLM.from_pretrained(model_path).half().cuda()
|