I think it is due to the following  -
torch.cuda.set_device(rank % torch.cuda.device_count())
Check the following:

case to execute model on specific gpu https://github.com/pytorch/examples/blob/master/imagenet/main.py#L144 15

case to execute on all available gpus
https://github.com/pytorch/examples/blob/master/imagenet/main.py#L154 26