Hi,
When I wrap my model in DataParallel my forward() method no longer gets input of the right size. For example
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, x, y):
        return x, y

net = Net()

device = torch.device('cuda')
net = nn.DataParallel(net)

net = net.to(device)

x = torch.ones((1024, 3, 100, 100))
y = torch.ones((1, 3, 100, 100))

x_out, y_out = net(x, y)

print(x_out.shape)
#torch.Size([128, 3, 100, 100])


I have 8 GPUs so it seems like I only get back the result from one of them.  If I change the first dimension of y to match x I get back the original result, as expected.