Rank 0 test on device cuda:0
Rank 1 test on device cuda:1
Test set: Average loss: 0.0274, Accuracy: 9913/10000 (99.13%)
RuntimeError: Expected tensor for argument #1 input to have
the same device as tensor for argument #2 weight;
but device 0 does not equal 1
(while checking arguments for cudnn_convolution)

This means the first parameter of both models are placed onto the correct device. Can you do the same check for all parameters? i.e., making sure that all parameters are placed to the correct device.

output = model(data)

Before the line above in test(...), can you print the device ids of the data as well? Looks like the model and data device does not match on rank 1.