Update README.md
Browse files
README.md
CHANGED
@@ -34,7 +34,7 @@ compute_dtype = torch.bfloat16
|
|
34 |
device = 'cuda'
|
35 |
model_id = "mobiuslabsgmbh/DeepSeek-R1-ReDistill-Qwen-1.5B-v1.1"
|
36 |
|
37 |
-
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=compute_dtype, attn_implementation=
|
38 |
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
39 |
|
40 |
chat = tokenizer.apply_chat_template([{"role":"user", "content":"What is 1.5+102.2?"}], tokenize=True, add_generation_prompt=True, return_tensors="pt")
|
|
|
34 |
device = 'cuda'
|
35 |
model_id = "mobiuslabsgmbh/DeepSeek-R1-ReDistill-Qwen-1.5B-v1.1"
|
36 |
|
37 |
+
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=compute_dtype, attn_implementation="sdpa", device_map=device)
|
38 |
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
39 |
|
40 |
chat = tokenizer.apply_chat_template([{"role":"user", "content":"What is 1.5+102.2?"}], tokenize=True, add_generation_prompt=True, return_tensors="pt")
|