Text Generation
Transformers
Safetensors
zamba2
Inference Endpoints
pglo commited on
Commit
019014a
·
verified ·
1 Parent(s): 51d1471

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
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="auto", torch_dtype=torch.bfloat16)
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")