Update README.md

README.md

@@ -25,11 +25,11 @@ model.lm_head = CastOutputToFloat(model.lm_head)
 
 tokenizer = transformers.AutoTokenizer.from_pretrained(base_path)
 
-batch = "<question>What are the symptoms of Alzheimer's disease?<answer>"
+batch = tokenizer.encode("<question>What are the symptoms of Alzheimer's disease?<answer>", return_tensors="pt")
 
 with torch.cuda.amp.autocast():
     out = model.generate(
-        input_ids=batch['input_ids'], 
+        input_ids=batch.to(model.device),
         max_length=300,
         do_sample=True,
         top_k=40,
@@ -38,6 +38,5 @@ with torch.cuda.amp.autocast():
         eos_token_id=tokenizer.additional_special_tokens_ids[tokenizer.additional_special_tokens.index('<question>')]
     )
 
-message = tokenizer.decode(out[0, :-1]).replace('<question>', "User:\n").replace('<answer>', 'Assistant:\n')
-
+print(tokenizer.decode(out[0, :-1]).replace('<question>', "User:\n").replace('<answer>', '\nAssistant:\n'))
 ```