RuntimeError: The size of tensor a (16) must match the size of tensor b (16000) at non-singleton dimension 2
### System Info

- `transformers` version: 4.30.0.dev0
- Platform: Linux-5.10.147+-x86_64-with-glibc2.31
- Python version: 3.10.11
- Huggingface_hub version: 0.14.1
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.0+cu118 (False)
- Tensorflow version (GPU?): 2.12.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.6.9 (cpu)
- Jax version: 0.4.8
- JaxLib version: 0.4.7
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: NO

### Who can help?

@ArthurZucker 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Run this notebook: https://colab.research.google.com/drive/1TFI84P9W4VPhNLgEngxPN57RwzS0C4bG?usp=sharing

### Expected behavior

Expected the model to train successfully.  Instead it gives a tensor mismatch error.
Hi @Tylersuard, thanks for reporting this issue. 

So that we can best try and help you, could you update the notebook so that it contains the minimal logic to replicate the error and can be run out-of-the-box? As it stands, there's many blocks with comments; references to loading / processing data we don't have access to;  doesn't currently have the reported error shown but does have many other errors. <|||||>Sorry @amyeroberts , Here is the updated version: https://colab.research.google.com/drive/1TFI84P9W4VPhNLgEngxPN57RwzS0C4bG?usp=sharing<|||||>I think you're splitting your input sequence into chunks of length 16: https://github.com/huggingface/transformers/blob/v4.29.1/src/transformers/models/mega/modeling_mega.py#L1063<|||||>@OllieBroadhurst That is correct.  As per the documentation (https://huggingface.co/docs/transformers/main/model_doc/mega) , I set the chunk_size equal to 16 and use_chunking to true, and the context length is a multiple of the chunk size.  My problem is not solved.<|||||>What I mean is have you tried turning chunking off?<|||||>@OllieBroadhurst Thank you for your suggestion.  I would likely run into out-of-memory errors, but I will try it.<|||||>Ok I tried it without chunking and I got out-of-memory errors.<|||||>This should still be adressed! Mega's forward pass might need some debugging. I can't do this fast, but keeping an eye on it! 