Repository: pytorch
Issue: Error with DataParallel and dataclass
Labels: triaged,module: data parallel

## 🐛 Bug

Something weird happens when mixing a DataParallel model with a dataclass input.


## To Reproduce

```python
from dataclasses import dataclass, asdict
import torch
from torch import nn

if torch.cuda.device_count() < 2:
    raise ValueError("Please test with mutiple GPUs")

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc = nn.Linear(2, 1)

    def forward(self, torch_sample):
        if isinstance(torch_sample, torch.Tensor):
            return self.fc(torch_sample)
        elif isinstance(torch_sample, dict):
            return self.fc(torch_sample["ft"])
        else:
            return self.fc(torch_sample.ft)

device = "cuda"
model = Model()
model = nn.DataParallel(model)
model = model.to(device)

@dataclass
class Sample:
    ft: torch.tensor = torch.tensor([])

        
ts = Sample(ft=torch.randn((10, 2)).cuda())
model(ts.ft)
model(asdict(ts))
model(ts)

ts = Sample(ft=torch.randn((10, 1, 2)).cuda())
model(ts.ft)
model(asdict(ts))
# With extra dim, it fails
model(ts)
```

It returns:

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)

     35 
     36 ts = Sample(ft=torch.randn((10, 1, 2)).cuda())
---> 37 model(ts)

...site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

.../site-packages/torch/nn/parallel/data_parallel.py in forward(self, *inputs, **kwargs)
    153             return self.module(*inputs[0], **kwargs[0])
    154         replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
--> 155         outputs = self.parallel_apply(replicas, inputs, kwargs)
    156         return self.gather(outputs, self.output_device)
    157 

.../site-packages/torch/nn/parallel/data_parallel.py in parallel_apply(self, replicas, inputs, kwargs)
    163 
    164     def parallel_apply(self, replicas, inputs, kwargs):
--> 165         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
    166 
    167     def gather(self, outputs, output_device):

.../site-packages/torch/nn/parallel/parallel_apply.py in parallel_apply(modules, inputs, kwargs_tup, devices)
     83         output = results[i]
     84         if isinstance(output, ExceptionWrapper):
---> 85             output.reraise()
     86         outputs.append(output)
     87     return outputs

...7/site-packages/torch/_utils.py in reraise(self)
    393             # (https://bugs.python.org/issue2651), so we work around it.
    394             msg = KeyErrorMessage(msg)
--> 395         raise self.exc_type(msg)

RuntimeError: Caught RuntimeError in replica 1 on device 1.
Original Traceback (most recent call last):
  File ".../site-packages/torch/nn/parallel/parallel_apply.py", line 60, in _worker
    output = module(*input, **kwargs)
  File ".../site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "...", line 15, in forward
    return self.fc(torch_sample.ft)
  File ".../site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "...s/torch/nn/modules/linear.py", line 91, in forward
    return F.linear(input, self.weight, self.bias)
  File ".../site-packages/torch/nn/functional.py", line 1678, in linear
    output += bias
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!
```

## Environment

PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: Could not collect
CMake version: version 3.15.3

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: Tesla P100-PCIE-16GB
GPU 1: Tesla P100-PCIE-16GB
GPU 2: Tesla P100-PCIE-16GB
GPU 3: Tesla P100-PCIE-16GB
GPU 4: Tesla P100-PCIE-16GB
GPU 5: Tesla P100-PCIE-16GB
GPU 6: Tesla P100-PCIE-16GB
GPU 7: Tesla P100-PCIE-16GB

Nvidia driver version: 450.66
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.16.2
[pip3] torch==1.0.0
[pip3] torchvision==0.2.1
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.0.130                      0  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py37he904b0f_0  
[conda] mkl_fft                   1.2.0            py37h23d657b_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] msgpack-numpy             0.4.7                      py_0  
[conda] numpy                     1.19.2                   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.2.0                    pypi_0    pypi
[conda] torch                     1.6.0                    pypi_0    pypi
[conda] torchsparseattn           0.3.dev0                  dev_0    <develop>
[conda] torchtext                 0.2.3                    pypi_0    pypi