Mismatch between input and target batch_sizes while training FSMT model
Author: HiGal
Pull request: no

Code to reproduce

```python
 tokenizer = get_fsmt_tokenizer()
    tokenizer.model_max_length=100
    model = get_fsmt_model()
    freeze_embeds(model)
    freeze_encoder(model)
    train_dataset = YandexRuEnDataset("data", split="train")
    val_dataset = YandexRuEnDataset("data", split="valid")

    training_args = TrainingArguments(
        output_dir='./results',  # output directory
        num_train_epochs=3,  # total # of training epochs
        per_device_train_batch_size=16,  # batch size per device during training
        per_device_eval_batch_size=64,  # batch size for evaluation
        warmup_steps=5000,  # number of warmup steps for learning rate scheduler
        weight_decay=0.01,  # strength of weight decay
        logging_dir='./logs',  # directory for storing logs
        fp16=True,
        fp16_opt_level='O2',
        save_steps=20000
    )

    trainer = Seq2SeqTrainer(
        model=model,  # the instantiated 🤗 Transformers model to be trained
        args=training_args,  # training arguments, defined above
        train_dataset=train_dataset,  # training dataset
        eval_dataset=val_dataset,  # evaluation dataset
        tokenizer=tokenizer,
        data_collator=collate_sentences(tokenizer)
    )

    trainer.train()
```
Dataset class

```python
class YandexRuEnDataset(Dataset):

    def __init__(self, root_data, split):
        src = open(f"{root_data}/corpus.en_ru.1m.ru", "r").readlines()
        tgt = open(f"{root_data}/corpus.en_ru.1m.en", "r").readlines()
        # src = src[:int(0.1*len(src))]
        # tgt = tgt[:int(0.1 * len(tgt))]
        X_train, X_test, y_train, y_test = train_test_split(src, tgt, test_size=0.33, random_state=228)
        if split == "train":
            self.src = X_train
            self.trg = y_train
        elif split == "valid":
            self.src = X_test
            self.trg = y_test

    def __len__(self):
        return len(self.src)

    def __getitem__(self, idx):
        src = self.src[idx]
        trg = self.trg[idx]

        return src, trg


def collate_sentences(tokenizer: Tokenizer):
    def collate_fn(batch):
        batch = list(zip(*batch))
        X_batch = list(batch[0])
        y_batch = list(batch[1])
        batch = tokenizer.prepare_seq2seq_batch(
            src_texts=X_batch, tgt_texts=y_batch,
            padding=True, truncation=True,
            return_tensors='pt'
        )
        return batch

    return collate_fn
```

Exception

```
  File "/home/farit/PycharmProjects/NMT/train.py", line 40, in <module>
    trainer.train()
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/transformers/trainer.py", line 940, in train
    tr_loss += self.training_step(model, inputs)
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/transformers/trainer.py", line 1302, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/transformers/trainer.py", line 1334, in compute_loss
    outputs = model(**inputs)
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/transformers/models/fsmt/modeling_fsmt.py", line 1180, in forward
    masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.tgt_vocab_size), labels.view(-1))
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 961, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/home/farit/anaconda3/envs/NMT/lib/python3.8/site-packages/torch/nn/functional.py", line 2261, in nll_loss
    raise ValueError('Expected input batch_size ({}) to match target batch_size ({}).'
ValueError: Expected input batch_size (1440) to match target batch_size (1600).
````

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.