Flos fix
<!-- This line specifies which issue to close after the pull request is merged. -->
Fixes #7146 

This basically unwraps the model that is used during training and can be either plain `Module` or `DataParallel`/`DistributedDataParallel`.
Please merge at will, as this fix is blocking us (https://github.com/huggingface/transformers/issues/7146#issuecomment-698852274). 