

TITLE
GPTNeo Error Attempting to Generate Text

COMMENTS
2

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cpu (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no
- Using TPU: Yes

### Who can help
@patil-suraj and @patrickvonplaten

Models:
GPTNeo

Library:

- flax transformers

 -->

## Information

Model I am using (Bert, XLNet ...): GPTNeo


## To reproduce

Steps to reproduce the behavior:

Here is a Google Colab for reproducing: https://colab.research.google.com/drive/1tba52h5t-BP3g13FMdPXVjKqpoLTlGvP?usp=sharing

For convenience here is the error msg:
```
TypeError: dynamic_update_slice update shape must be smaller than operand shape, got update shape (1, 45) for operand shape (1, 20).
```

I was originally getting the same error as #12081. However, when I attempted to implement the same fix as in that issue, I got the above error. The error might be because I am using the "ForCausalLM" version of GPTNeo. However, there is no LMHead version

## Expected behavior
Generate the output sequence
