Properly set LightningModule.device after model replacement (#7188)

Co-authored-by: Carlos Mocholí <[email protected]>

CHANGELOG.md

@@ -325,6 +325,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed parsing for pre-release package versions ([#6999](https://github.com/PyTorchLightning/pytorch-lightning/pull/6999))
 
 
+- Fixed resetting device after `fitting/evaluating/predicting` ([#7188](https://github.com/PyTorchLightning/pytorch-lightning/pull/7188))
+
+
 ## [1.2.7] - 2021-04-06
 
 ### Fixed

pytorch_lightning/accelerators/gpu.py

@@ -47,9 +47,10 @@ def on_train_start(self) -> None:
         with torch.cuda.device(self.root_device):
             torch.cuda.empty_cache()
 
-    def on_train_end(self) -> None:
+    def teardown(self) -> None:
+        self.lightning_module.cpu()
+
         # clean up memory
-        self.model.cpu()
         with torch.cuda.device(self.root_device):
             torch.cuda.empty_cache()
 

pytorch_lightning/utilities/device_dtype_mixin.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from typing import Optional, Union
-
 import torch
 from torch.nn import Module
 
@@ -47,11 +46,6 @@ def device(self) -> Union[str, torch.device]:
 
         return device
 
-    @device.setter
-    def device(self, new_device: Union[str, torch.device]):
-        # Necessary to avoid infinite recursion
-        raise RuntimeError('Cannot set the device explicitly. Please use module.to(new_device).')
-
     @parameter_validation
     def to(self, *args, **kwargs) -> Module:
         """Moves and/or casts the parameters and buffers.

tests/core/test_lightning_module.py

@@ -14,13 +14,15 @@
 from unittest.mock import Mock
 
 import pytest
+import torch
 from torch import nn
 from torch.optim import Adam, SGD
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
+from tests.helpers.runif import RunIf
 
 
 def test_property_current_epoch():
@@ -358,3 +360,24 @@ def configure_optimizers(self):
     )
 
     trainer.fit(model)
+
+
+@RunIf(min_gpus=1)
+def test_device_placement(tmpdir):
+
+    model = BoringModel()
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, gpus=1)
+    trainer.fit(model)
+
+    def assert_device(device: torch.device) -> None:
+        assert model.device == device
+        for p in model.parameters():
+            assert p.device == device
+
+    assert_device(torch.device("cpu"))
+    model.to(torch.device("cuda:0"))
+    assert_device(torch.device("cuda:0"))
+    trainer.test(model)
+    assert_device(torch.device("cpu"))
+    trainer.predict(model, dataloaders=model.train_dataloader())
+    assert_device(torch.device("cpu"))