HI, I try following code with multi gpus, but the command “loss.backward()” throws an error “RuntimeError: arguments are located on different GPUs. The pytorch version is 1.0.0. How can i do it?
from torch import nn
import torch

class MyModule1(nn.Module):
  def __init__(self):
    super(MyModule1, self).__init__()
    self.fc1 = nn.Linear(2, 1)

  def forward(self, x):
    x = self.fc1(x)
    return x

class MyModule2(nn.Module):
  def __init__(self):
    super(MyModule2, self).__init__()
    self.fc2 = nn.Linear(1, 1)

  def forward(self, x):
    x = self.fc2(x)
    return x

model1 = MyModule1()
model1.to('cuda:0')
x = torch.tensor([10, 5], dtype=torch.float).to('cuda:0')
y = model1(x)
z = y.to('cuda:1')
model2 = MyModule2()
model2.to('cuda:1')
w = model2(y)
w.backward()