Error In Fine-Tuning MarianMT's opus-mt model ValueError: The two structures don't have the same sequence length. Input structure has length 4, while shallow structure has length 3
## Environment info
- `transformers` version: 4.11.3
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.12
- PyTorch version (GPU?): 1.9.0+cu111 (False)
- Tensorflow version (GPU?): 2.6.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes 
- Using distributed or parallel set-up in script?: No 

### Who can help
@patrickvonplaten 

## Information

Model I am using MarianMT:

The problem arises when using:

* [ ] my own modified scripts: I aim to finetune the opus-mt-en-hi model on a custom dataset for En-Hindi translation and while using the cross entropy objective my model.fit() loop errors out.  
I am facing the same error as mentioned in #11560 . In my best understanding, the source of the error is that the shape of y_pred and y_true don't match such as - y_pred is (Batch_size, seq_length, vocab_size ) and y_true is ( Batch_size, seq_length) but I am not sure if I have to softmax ( to generate y_pred per element that can be computed against each element in y_true ) the logits first to get the training started. 

The tasks I am working on is:
* task: translation

## To reproduce : 
End-End script available [here](https://colab.research.google.com/drive/11wVlquduPyCV7vb6nh1eEt2PhvAbkyBM?usp=sharing) 

Steps to reproduce the behavior:

Run colab script ( sample tensors are already provided as inputs) 

## Expected behavior

I looked at
https://github.com/huggingface/transformers/blob/master/examples/tensorflow/translation/run_translation.py but except the custom objective function ( which is adding functionality to the same base sparse CE function ) , cannot identify anything that could be altered in my current script. Hence, the expected behavior would be for the training to proceed as expected. 

Would be very helpful if you could provide some direction on how to resolve this.
Thanks in advance!

This seems related to TF model, gently pinging @Rocketknight1 :) <|||||>@Rocketknight1 - lemme know if you think you'll find time. Otherwise I can take a look as well I think :-)<|||||>I'll take a look!<|||||>@patrickvonplaten Investigated for an hour or so - I was able to reproduce the error, but debugging the script ran into constant shape errors no matter what I tried. I think I'd need to dig into this in more detail and make sure the input data is really what the model is expecting, but I have a bunch of other priorities right now, so I probably won't get to it in the next few days at least!<|||||>Sounds good! Let me know if you would like me to take over the issue<|||||>Hi, I'm sorry for how long it took me to get to this. Once I actually had time to sit down and look at it, the bug is straightforward - it's not within the model at all, but within the data preparation. The TF dataset you're creating does not batch inputs, and therefore the first batch dimension is missing, which is what causes the downstream shape errors. There are secondary issues with the loss function as supplied - as the model outputs logits, SparseCategoricalCrossentropy() should be called only with `from_logits=True`. You may find it easier to use the internal loss by leaving the `loss` argument blank.

Here's [a modified version of your script](https://colab.research.google.com/drive/11wVlquduPyCV7vb6nh1eEt2PhvAbkyBM?usp=sharing) that should work. Please note, though, that we can't do extensive bug-fixing of user code in GitHub issues - taking this to the Discord or forums would usually be better!<|||||>> Hi, I'm sorry for how long it took me to get to this. Once I actually had time to sit down and look at it, the bug is straightforward - it's not within the model at all, but within the data preparation. The TF dataset you're creating does not batch inputs, and therefore the first batch dimension is missing, which is what causes the downstream shape errors. There are secondary issues with the loss function as supplied - as the model outputs logits, SparseCategoricalCrossentropy() should be called only with `from_logits=True`. You may find it easier to use the internal loss by leaving the `loss` argument blank.
> 
> Here's [a modified version of your script](https://colab.research.google.com/drive/11wVlquduPyCV7vb6nh1eEt2PhvAbkyBM?usp=sharing) that should work. Please note, though, that we can't do extensive bug-fixing of user code in GitHub issues - taking this to the Discord or forums would usually be better!

Thanks a lot for looking into it! Noted.