Update README.md

README.md

@@ -39,6 +39,13 @@ past_key_values = outputs.past_key_values
 # next step
 outputs = model(input_ids=input_ids, use_cache=True, past_key_values=past_key_values)
 
-######## Generate?
+######## Generate!
 with torch.autocast(**amp_settings), torch.no_grad():
-    model.generate(input_ids, max_new_tokens=20, use_cache=True)
+    output_ids = model.generate(input_ids, max_new_tokens=20, use_cache=True, num_steps=32)
+print(tokenizer.decode(output_ids[0]))
+
+
+# with or without cache
+with torch.autocast(**amp_settings), torch.no_grad():
+    output_ids = model.generate(input_ids, max_new_tokens=20, use_cache=False, num_steps=32)
+print(tokenizer.decode(output_ids[0]))