jon-tow commited on
Commit
8dd6378
1 Parent(s): ae8a1b5

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -2
README.md CHANGED
@@ -21,9 +21,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
21
 
22
  tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-base-alpha-7b")
23
  model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-base-alpha-7b")
24
- model.half()
25
 
26
- inputs = tokenizer("What's your mood today?", return_tensors="pt")
27
  tokens = model.generate(
28
  **inputs,
29
  max_new_tokens=64,
 
21
 
22
  tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-base-alpha-7b")
23
  model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-base-alpha-7b")
24
+ model.half().cuda()
25
 
26
+ inputs = tokenizer("What's your mood today?", return_tensors="pt").to("cuda")
27
  tokens = model.generate(
28
  **inputs,
29
  max_new_tokens=64,