What pytorch version are you using? Pinning to specific GPU was merged recently https://github.com/pytorch/pytorch/pull/4196, before that context was always initialized on device 0.
What pytorch version are you using? Pinning to specific GPU was merged recently https://github.com/pytorch/pytorch/pull/4196, before that context was always initialized on device 0.