class Dummy(nn.Module):
    def __init__(self, input_shape):
        super(Dummy, self).__init__()
        self.parameter = nn.Parameter(torch.randn(*input_shape), requires_grad=True)

    def forward(self, x):
        return x + self.parameter
        

model = Dummy((1, 2, 3)).to('cuda:0')
input = torch.randn((1, 2, 3))
input = input.to('cuda:0')
output = model(input)
output = output.to('cpu')

model.to('cpu')
loss_fn = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)

loss = loss_fn(output, golden_output)

optimizer.zero_grad()
loss.backward()
optimizer.step()

I know the above example is pretty useless but since the module has no intermediate outputs, when the output + parameters are moved to cpu wouldn’t the backwards pass be able to run on CPU.
In this way, if this was a larger module with more parameters and gradient outputs that are moved to a different device, wouldn’t the backwards pass be able to run there?
From my understanding of Loss.backward() throws an error with multi gpus 5 as long as the tensors are moved to the appropriate device the operations should be able to run.