Update README.md
Browse files
README.md
CHANGED
@@ -35,7 +35,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
35 |
import torch
|
36 |
|
37 |
tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
|
38 |
-
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-2.7B", device_map="
|
39 |
|
40 |
input_text = "A funny prompt would be "
|
41 |
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
|
|
|
35 |
import torch
|
36 |
|
37 |
tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
|
38 |
+
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-2.7B", device_map="cuda", torch_dtype=torch.bfloat16)
|
39 |
|
40 |
input_text = "A funny prompt would be "
|
41 |
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
|