Update README.md

README.md

@@ -47,9 +47,9 @@ model = BloomForCausalLM.from_pretrained(ckpt).to(device)
 
 def generate_prompt(text):
     inputs = tokenizer(text, return_tensors='pt')
-    input_ids = inputs.input_ids.to("cuda")
-    attention_mask = inputs.attention_mask.to("cuda")
-    output = model.generate(input_ids, attention_mask=attention_mask, max_length=2048, eos_token_id=tokenizer.eos_token_id) # num_beams=3, temperature=1.9
+    input_ids = inputs.input_ids.to(device)
+    attention_mask = inputs.attention_mask.to(device)
+    output = model.generate(input_ids, attention_mask=attention_mask, max_length=2048, eos_token_id=tokenizer.eos_token_id)
 
     return tokenizer.decode(output[0], skip_special_tokens=False)