set device to root gpu in single GPU case

[ananyahjha93]

Fixes #3030

are there suggested tests for this PR?

[nateraw]

LGTM 😄

If possible, mind adding quick test for this? Approving so we don't block this bugfix either way.

[ananyahjha93]

@awaelchli weren't the random failing tests fixed?

[awaelchli]

@ananyahjha93 I think a simple test would be to check
torch.cuda.current_device() == self.root_gpu
for the relevant backends, if I understand correctly.

[awaelchli]

@ananyahjha93 I thought so. Can we rerun the CI to check if it's the same accuracy value. I want to see if it's deterministic.
One thing is suspicious is that in the test we set os.environ['CUDA_VISIBLE_DEVICES'] = '0,1', and this PR is setting the device also.

[ananyahjha93]

@awaelchli should pass the GPU tests, the failure was ddp_spawn test on the accuracy metric.

[awaelchli]

the test passes now in a second run, which means despite setting the deterministic flag in the Trainer and seed_everything(1234), something in that test is still random.

[codecov]

Codecov Report

Merging #3042 into master will decrease coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff           @@
##           master   #3042   +/-   ##
======================================
- Coverage      90%     90%   -0%     
======================================
  Files          81      81           
  Lines        7672    7685   +13     
======================================
+ Hits         6906    6912    +6     
- Misses        766     773    +7     

[Borda]

still no test, right?