make device property always return a device with index (#4851)

* make device property always return a device with index

* pep8

* Update test_dtype_device_mixin.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

pytorch_lightning/utilities/device_dtype_mixin.py

@@ -37,7 +37,13 @@ def dtype(self, new_dtype: Union[str, torch.dtype]):
 
     @property
     def device(self) -> Union[str, torch.device]:
-        return self._device
+        device = self._device
+
+        # make this more explicit to always include the index
+        if device.type == 'cuda' and device.index is None:
+            return torch.device(f'cuda:{torch.cuda.current_device()}')
+
+        return device
 
     @device.setter
     def device(self, new_device: Union[str, torch.device]):

tests/utilities/test_dtype_device_mixin.py

@@ -55,7 +55,6 @@ def on_train_batch_start(self, trainer, model, batch, batch_idx, dataloader_idx)
 ])
 @pytest.mark.parametrize(['dst_device'], [
     pytest.param(torch.device('cpu')),
-    pytest.param(torch.device('cuda')),
     pytest.param(torch.device('cuda', 0)),
 ])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
@@ -100,3 +99,16 @@ def test_submodules_multi_gpu_ddp_spawn(tmpdir):
         max_steps=1,
     )
     trainer.fit(model)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_gpu_device_includes_index():
+    model = TopModule()
+
+    # explicitly call without an index to see if the returning device contains an index (it should!)
+    model.cuda()
+
+    device = model.device
+    assert device.type == 'cuda'
+    assert device.index is not None
+    assert device.index == torch.cuda.current_device()