Hi, I have the same problem as you. I want to train the same model in two different GPU.
import torch
import torch.nn as nn
import torch.optim as optim

class ToyModel(nn.Module):
  def __init__(self):
    super(ToyModel, self).__init__()
    self.net1 = torch.nn.Linear(10, 10)  
    self.relu = torch.nn.ReLU()
    self.net2 = torch.nn.Linear(10, 5)  

  def forward(self, x):
    x = self.relu(self.net1(x))
    return self.net2(x)
    
model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()

outputs1 = model.cuda(1)(torch.randn(20, 10).cuda(1)) #gpu=1
outputs2 = model.cuda(2)(torch.randn(20,10).cuda(2)) #gpu=2
labels = torch.randn(20, 5).cuda(3) #gpu=3 loss
loss_fn(outputs1.cuda(3)+outputs2.cuda(3),labels).backward()
optimizer.step()

But the error is RuntimeError: binary_op(): expected both inputs to be on same device, but input a is on cuda:1 and in│
put b is on cuda:2