So after a few weeks of experimenting with this functionality, we have discovered this whole thing is a little bit more complicated than we initially thought.
As mentioned above in this thread, there is a problem with torch.nn.DataParallel. It also causes problem when exporting to ONNX or others. Those can be fixed, but the biggest problem is that TensorFlow does not like those new model output types at all, so we are now left with two different APIs for PyTorch and TensorFlow, which is not what we want for the library.
On the other hand, we have a version that works with TensorFlow and doesn’t break torch.nn.DataParallel in #5981 7. Its problem is that it breaks backward-compatibility since this model output behaves more like a dict than tuple, and unpacking it (e.g. doing loss, logits = output) iters through the keys and not the values. This breaks pretty much all tests and examples since we use that unpacking a lot.
After lots of internal discussion, and since we can’t find a way to have the best of both world (e.g., fix the current model output without breaking change to work with TensorFlow and torch.nn.DataParallel) we have decided to revert a bit those changes. I’ll start implementing the following tomorrow:

The current ModelOutput will become the more-dict-version in #5981 7 and it will be used for all TensorFlow models as well as PyTorch.
Since this is a breaking change, we’ll switch the flip on the other side: instead of getting the new model outputs and having the possibility of getting back a tuple with return_tuple=True, we will go back to return tuple by default and have a return_dict argument you can pass to get the new model outputs.
Update all our examples to use return_dict=True and discourage the use of unpacking, so that everyone has time to get use to the future breaking change to come. Test and fix potential side effects.
In a few months (and probably during a major release), switch the flip and make return_dict=True the default for all models