If statement occurs an error under the multi-gpu environment

Hi, I used this model, https://github.com/ozan-oktay/Attention-Gated-Networks, and I found an error,
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)
The author of this model uses only one gpu, but we have 4 as well. (4 V100 with NVLink)
Also when we remove if statement or use only one gpu, it works well.
What’s the problem? I need your help. Thanks.


@Hyungjoo_Cho I encounter the same problem, but not in the same code. Have you managed to sort it out?

torch.cuda.set_device(1), you can use it on the top of codes