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

autograd
1

what I want to do is like this:

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

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

loss.backward()

loss.backward() will work?

2

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