FYI, this new feature breaks when using nn.DataParallel (multi gpu) as currently torch.nn.parallel.scatter_gather.gather can’t gather outputs that are dataclasses  6.
There are workarounds worked on in the core, but until then if you run into this issue in your custom code you may need to add:
+            if isinstance(model, torch.nn.DataParallel):
+                inputs["return_tuple"] = True

before your call:
             outputs = model(**inputs)

For more details see: https://github.com/huggingface/transformers/issues/5693 9
p.s. voting on the pytorch issue 6 may help resolve this globally faster.