Could you try to get the current device from the passed tensors instead of
int64_t curDevice = at::cuda::current_device();  

I haven’t tested the code yet, but if I’m not mistaken, this would use the current device specified by a device guard.