How to do forward step on one gpu and calculate loss on another gpu?

what I want to do is like this:

# on cuda:0
output = model.forward(input)

# on cuda:1
# calculate perceptron loss
loss = perceptron_loss(output, target)


loss.backward() will work?

It’s not a fast way of doing it. You’d have to copy the information on gpu0 -> cpu -> gpu1 then do backward. And maybe the model also have to be copied like that. At least to my knowledge