Repository: pytorch
Issue: Why does using nn.DataParallel add a dimension after the nn.Parameter parameter? Then it results in a dimension mismatch and an error when multiplying.
Labels: triaged,module: data parallel

### 🐛 Describe the bug

class RFNO(nn.Module):
    def __init__(self, out_channels, modes1, modes2):
        super(RFNO, self).__init__()

        self.out_channels = out_channels
        self.modes1 = modes1
        self.modes2 = modes2

        self.scale = (1 / out_channels)
        self.weights0 = self.scale * torch.rand(1, out_channels, 1, 1, dtype=COMPLEXTYPE)
        self.weights1 = self.scale * torch.rand(1, out_channels, self.modes1, self.modes2, dtype=COMPLEXTYPE)
        self.weights2 = self.scale * torch.rand(1, out_channels, self.modes1, self.modes2, dtype=COMPLEXTYPE)

        self.weights0 = nn.Parameter(self.weights0)
        self.weights1 = nn.Parameter(self.weights1)
        self.weights2 = nn.Parameter(self.weights2)

    # Complex multiplication
    def compl_mul2d(self, input, weights):
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x):
        batchsize = x.shape[0]
        
        # Move weights to the same device as input `x`
        weights0 = self.weights0.to(x.device)
        weights1 = self.weights1.to(x.device)
        weights2 = self.weights2.to(x.device)

        x_ft = torch.fft.rfft2(x)
        x_ft = x_ft * weights0
        
        # Multiply relevant Fourier modes
        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1)//2 + 1, dtype=COMPLEXTYPE, device=x.device)
        out_ft[:, :, :self.modes1, :self.modes2] = self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], weights2)

        # Return to physical space
        print(out_ft.shape)
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x

### Versions

When I'm not using nn.DataParallel, the shape of each parameter is as follows：
#weights0.shape, weights1.shape, weights2.shape =  torch.Size([1, 64, 1, 1]) torch.Size([1, 64, 32, 32]) torch.Size([1, 64, 32, 32])
# x.shape = torch.Size([32, 1, 320, 320])
# x_ft.shape =  torch.Size([32, 1, 320, 161])

When I'm using nn.DataParallel, the shape of each parameter is as follows：
#weights0.shape, weights1.shape, weights2.shape =  torch.Size([1, 64, 1, 1,2]) torch.Size([1, 64, 32, 32,2]) torch.Size([1, 64, 32, 32,2])
# x.shape = torch.Size([32, 1, 320, 320])
# x_ft.shape =  torch.Size([32, 1, 320, 161])

Then it results in a dimension mismatch and an error when multiplying：
Original Traceback (most recent call last):
  File "/home/user/anaconda3/envs/opc/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in _worker
    output = module(*input, **kwargs)
  File "/home/user/anaconda3/envs/opc/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/anaconda3/envs/opc/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/opc/lithobench-DP/lithobench/litho/doinn.py", line 168, in forward
    br0 = self.rfno(F.avg_pool2d(x, kernel_size=8, stride=8))
  File "/home/user/anaconda3/envs/opc/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/anaconda3/envs/opc/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/opc/lithobench-DP/lithobench/litho/doinn.py", line 124, in forward
    x_ft = x_ft * weights0
RuntimeError: The size of tensor a (161) must match the size of tensor b (2) at non-singleton dimension 4