Update README.md

README.md

@@ -13,12 +13,17 @@ Bloom (2.5 B) Scientific Model fine-tuned on Zen knowledge
 
 
 ```python
+
+
+cuda = False
+
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("MultiTrickFox/bloom-2b5_Zen") 
 model = AutoModelForCausalLM.from_pretrained("MultiTrickFox/bloom-2b5_Zen")
 
-model
+model.cuda() if cuda else None
 tokenizer.pad_token_id = tokenizer.eos_token_id
 
 generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
@@ -26,7 +31,7 @@ generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
 inp = [ """Today""", """Yesterday""" ]
 
 out = generator( 
-    inp, do_sample=True,
+    inp.cuda() if cuda else inp, do_sample=True,
 
     temperature=.6,
     typical_p=.7,