Hello,
A model has time-series structure.
So, I’ve split each time-step into each gpu, due to insufficiency of gpu memory.
But I got an error when calling loss.backward():
File "/home/xxx/.local/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function CatBackward returned an invalid gradient at index 1 - expected device cuda:2 but got cuda:0

Please find the minimum code below:
    self.config['DEVICE_ALL'] = [0, 2, 3]
    self.backbone = UNet(n_channels=config['IN_LEN']*3 - 2)
    self.geo = nn.Parameter(data=torch.randn(1, 4, self.h, self.w), requires_grad=True)
    self.outConv = nn.Conv2d(68, 1, kernel_size=1)
    self.output_time_length = 3

    def forward(self, input):

        outputs = []
        cur_input = input
        cur_state = self.get_state(...)

        for i in range(self.output_time_length):
            
            dev = self.config['DEVICE_ALL'][i]
            cur_input = cur_input.cuda(dev)
            cur_state = cur_state.cuda(dev)
            x = torch.cat([cur_input, cur_state], 1).cuda(dev)

            self.backbone = self.backbone.cuda(dev)
            x = self.backbone(x).cuda(dev)

            geo = self.geo.expand(x.shape[0], -1, -1, -1).cuda(dev)
            x = torch.cat([x, geo], 1).cuda(dev)

            self.outConv = self.outConv.cuda(dev)
            x = self.outConv(x).cuda(dev)

            outputs.append(x)
            cur_input = torch.cat([cur_input[:, 1:], x], 1)
            cur_state = self.get_next_state(...)

        return torch.cat(outputs, 1).cuda(dev)

Thanks in advance!