UniT Training Tensor Size Mismatch

## 🐛 Bug

I am attempting to run training on the UniT single_task model. However I am getting a tensor size mismatch.

## Command

CUDA_VISIBLE_DEVICES=0 python mmf_cli/run.py config=projects/unit/configs/coco/single_task.yaml datasets=detection_coco model=unit run_type=train training.batch_size=8 checkpoint.resume_zoo=unit.coco.single_task

## To Reproduce

This occurs on a fresh clone of the repository.

The full error message:

Traceback (most recent call last):
File "mmf_cli/run.py", line 137, in
run()
File "mmf_cli/run.py", line 133, in run
main(configuration, predict=predict)
File "mmf_cli/run.py", line 56, in main
trainer.train()
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/trainers/mmf_trainer.py", line 142, in train
self.training_loop()
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/trainers/core/training_loop.py", line 33, in training_loop
self.run_training_epoch()
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/trainers/core/training_loop.py", line 77, in run_training_epoch
for idx, batch in enumerate(self.train_loader):
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/datasets/multi_dataset_loader.py", line 199, in __next__
next_batch = next(self.current_iterator)
File "/home/ahavrilla/miniconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 517, in __next__
data = self._next_data()
File "/home/ahavrilla/miniconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1199, in _next_data
return self._process_data(data)
File "/home/ahavrilla/miniconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1225, in _process_data
data.reraise()
File "/home/ahavrilla/miniconda3/lib/python3.8/site-packages/torch/_utils.py", line 429, in reraise
raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
File "/home/ahavrilla/miniconda3/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 202, in _worker_loop
data = fetcher.fetch(index)
File "/home/ahavrilla/miniconda3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
return self.collate_fn(data)
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/common/batch_collator.py", line 11, in __call__
sample_list = convert_batch_to_sample_list(batch)
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/common/sample.py", line 413, in convert_batch_to_sample_list
sample_list = SampleList(batch)
File "/home/ahavrilla/zensors-inc/mmf_vanilla/mmf/mmf/common/sample.py", line 130, in __init__
self[field][idx] = self._get_data_copy(sample[field])
RuntimeError: The expanded size of the tensor (704) must match the existing size (768) at non-singleton dimension 2. Target sizes: [3, 938, 704]. Tensor sizes: [3, 1038, 768]

## Expected behavior

Training

## Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

Collecting environment information...
PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: P106-100
Nvidia driver version: 460.73.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] pytorch-lightning==1.4.0.dev0
[pip3] torch==1.8.1
[pip3] torchaudio==0.8.1
[pip3] torchmetrics==0.3.2
[pip3] torchtext==0.5.0
[pip3] torchvision==0.9.1
[conda] numpy 1.20.3 pypi_0 pypi
[conda] pytorch-lightning 1.4.0.dev0 pypi_0 pypi
[conda] torch 1.8.1 pypi_0 pypi
[conda] torchaudio 0.8.1 pypi_0 pypi
[conda] torchmetrics 0.3.2 pypi_0 pypi
[conda] torchtext 0.5.0 pypi_0 pypi
[conda] torchvision 0.9.1 pypi_0 pyp

You can run the script with:
```
# For security purposes, please check the contents of collect_env.py before running it.
python -m torch.utils.collect_env
```

- PyTorch Version (e.g., 1.0): 1.8.1+cu102
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.8.5
- CUDA/cuDNN version: V9.1.85
- GPU models and configuration: 1 GPU
- Any other relevant information:

## Additional context

Reproduce the command from projects/unit/configs/coco/single_task.yaml and trace the batch shapes through mmf/common/batch_collator.py and mmf/common/sample.py. Investigate why tensors sized [3, 1038, 768] are copied into [3, 938, 704], then verify that UniT single-task training starts without the DataLoader size-mismatch error.

Written by the indexing model from the issue text.