problem in the doc, in the  "Quick Start" GPT2 example
I am going through the GPT2 example in the doc. Is there a mistake in the "Using the past" code. The main loop to generate text is:

```python
for i in range(100):
    print(i)
    output, past = model(context, past=past)
    token = torch.argmax(output[0, :])
    generated += [token.tolist()]
    context = token.unsqueeze(0)
```


At the first iteration the tensor output as shape [1, 3, 50257], and at all the following iterations it has size [1,50257]. Should the code be:

```python
for i in range(100):
    print(i)
    output, past = model(context, past=past)
    if i==0:
        token = torch.argmax(output[0,-1,:])
    else:
        token = torch.argmax(output[0, :])
    generated += [token.tolist()]
    context = token.unsqueeze(0)
```






This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
