Correct quickstart example when using the past

fd639e5b · Lysandre · 63a5399b · fd639e5b
Commit fd639e5b authored Feb 10, 2020 by Lysandre
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

docs/source/quickstart.md docs/source/quickstart.md +1 -1

No files found.
--- a/docs/source/quickstart.md
+++ b/docs/source/quickstart.md
@@ -209,7 +209,7 @@ past = None
 for i in range(100):
    print(i)
    output, past = model(context, past=past)
-    token = torch.argmax(output[0, :])
+    token = torch.argmax(output[..., -1, :])

    generated += [token.tolist()]
    context = token.unsqueeze(0)