I have read posts about model sharding and have found really good examples.
However, when I implemented it, at the 2nd epoch, I got the following error:
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 0 does not equal 1 (while checking arguments for cudnn_convolution)

My code is like:
class MyModel(nn.Module):
    def __init__(self, some_param):
        self.large_submodule1 = SubModel1(...)
        self.large_submodule2 = SubModel2(...)

        self.large_submodule1.cuda(0)
        self.large_submodule2.cuda(1)

    def forward(self, x):
        x1, x2, cat = self.large_submodule1(x)
        x1 = x1.cuda(1)
        x2 = x2.cuda(1)
        cat = cat.cuda(1)
        out = self.large_submodule2(x1, x2, cat)
        return out.cuda(0) # because the ground truth is saved at cuda(0)

class SubModel1(nn.Module):
    def __init__(self, some_param):
        self.conv1 = ...
        self.conv2 = ...
        self.sigmoid = 

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        cat = torch.cat((x1, x2), dim=1)
        return x1, x2, cat

class SubModel2(nn.Module):
    def __init__(self, some_param):
        self.conv1 = ...
        self.conv2 = ...
        self.sigmoid = 

    def forward(self, x1, x2, cat):
        x1 = self.conv1(x1)
        x2 = self.conv2(x2)
        out = torch.cat((x1, x2, cat), dim=1)
        return out

It worked at the 1st epoch, but throwed the above error at the 2nd epoch.
I’ve read this issue 2 about a similar error, but I didn’t use DataParallel and my batch size was 1.
Thank you for your time!