Repository: pytorch
Issue: torch.nn.parallel.scatter_gather.gather can't gather outputs that are dataclasses
Labels: oncall: distributed,triaged,module: data parallel

## 🐛 Bug

`torch.nn.parallel.scatter_gather.gather` can't gather outputs that are `dataclasses`

## To Reproduce

Steps to reproduce the behavior:

```
from dataclasses import dataclass

# base model outputs class in transformers
class ModelOutput:
    def to_tuple(self):
        return tuple(getattr(self, f) for f in self.__dataclass_fields__.keys() if getattr(self, f, None) is not None)

    def to_dict(self):
        return {f: getattr(self, f) for f in self.__dataclass_fields__.keys() if getattr(self, f, None) is not None}

    def __getitem__(self, i):
        return self.to_dict()[i] if isinstance(i, str) else self.to_tuple()[i]

    def __len__(self):
        return len(self.to_tuple())

@dataclass
class X1(ModelOutput):
    foo: torch.FloatTensor
    bar: torch.FloatTensor

x1 = X1(torch.tensor(1., device="cuda:0"), torch.tensor(2., device="cuda:0"))
x2 = X1(torch.tensor(3., device="cuda:1"), torch.tensor(4., device="cuda:1"))
outputs = [x1, x2]

from torch.nn.parallel.scatter_gather import gather
gather(outputs, 0)
```

Error:
```
TypeError: __init__() missing 1 required positional argument: 'bar'
```

as it fails to detect dataclass as dict, it skips over it and calls the class with wrong arguments, sending just the first argument.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

Should handle `dataclass` `outputs` as `dict` transparently

The following implementation fixes this simple example (but see later that it doesn't always work). 

Not sure whether pytorch's requirements include `dataclasses`:

```
import dataclasses
def gather(outputs, target_device, dim=0):
    r"""
    Gathers tensors from different GPUs on a specified device
      (-1 means the CPU).
    """
    def gather_map(outputs):
        out = outputs[0]
        if dataclasses.is_dataclass(out):
            outputs = [dataclasses.asdict(out) for out in outputs]
            out = outputs[0]
        if isinstance(out, torch.Tensor):
            return Gather.apply(target_device, dim, *outputs)
        if out is None:
            return None

        if isinstance(out, dict):
            if not all((len(out) == len(d) for d in outputs)):
                raise ValueError('All dicts must have the same number of keys')
            return type(out)(((k, gather_map([d[k] for d in outputs]))
                              for k in out))

        return type(out)(map(gather_map, zip(*outputs)))

    # Recursive function calls like this create reference cycles.
    # Setting the function to None clears the refcycle.
    try:
        res = gather_map(outputs)
    finally:
        gather_map = None
    return res
```

The lack of dataclasses outputs support was detected in transformers: https://github.com/huggingface/transformers/issues/5693#issuecomment-657325298

I saw that named tuples also have this issue: https://github.com/pytorch/vision/issues/1048



## Environment

```
Collecting environment information...
PyTorch version: 1.5.1
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration:
GPU 0: GeForce GTX TITAN X
GPU 1: GeForce GTX TITAN X

Nvidia driver version: 440.95.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.13.3
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] mkl                       2020.0                      166
[conda] mkl-service               2.3.0            py37he904b0f_0
[conda] mkl_fft                   1.0.15           py37ha843d7b_0
[conda] mkl_random                1.1.0            py37hd6b4f25_0
[conda] numpy                     1.18.1                   pypi_0    pypi
[conda] numpy-base                1.18.1           py37hde5b4d6_1
[conda] numpydoc                  0.9.2                      py_0
[conda] pytorch                   1.5.1           py3.7_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] pytorch-lightning         0.7.6                    pypi_0    pypi
[conda] pytorch-nlp               0.5.0                    pypi_0    pypi
[conda] pytorch-pretrained-bert   0.6.2                    pypi_0    pypi
[conda] pytorch-transformers      1.1.0                    pypi_0    pypi
[conda] torch                     1.5.0                    pypi_0    pypi
[conda] torchtext                 0.5.1                    pypi_0    pypi
[conda] torchvision               0.6.1                py37_cu102    pytorch

```

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528 @osalpekar @jiayisuse @agolynski