Add a bit of a hack to fix self._device

Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Nov 15, 2024 · 86de71e · 86de71e
1 parent d7cef8b
commit 86de71e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/project/algorithms/testsuites/algorithm_tests.py b/project/algorithms/testsuites/algorithm_tests.py
@@ -65,6 +65,9 @@ def test_initialization_is_deterministic(
             assert isinstance(algorithm_1, lightning.LightningModule)
 
             with trainer.init_module():
+                # A bit hacky, but we have to do this because the lightningmodule isn't associated
+                # with a Trainer.
+                algorithm_1._device = torch.get_default_device()
                 algorithm_1.configure_model()
 
         with torch.random.fork_rng(devices=list(range(torch.cuda.device_count()))):
@@ -73,6 +76,9 @@ def test_initialization_is_deterministic(
             assert isinstance(algorithm_2, lightning.LightningModule)
 
             with trainer.init_module():
+                # A bit hacky, but we have to do this because the lightningmodule isn't associated
+                # with a Trainer.
+                algorithm_2._device = torch.get_default_device()
                 algorithm_2.configure_model()
 
         torch.testing.assert_close(algorithm_1.state_dict(), algorithm_2.state_dict())
@@ -157,6 +163,9 @@ def test_initialization_is_reproducible(
             algorithm = instantiate_algorithm(experiment_config.algorithm, datamodule=datamodule)
             assert isinstance(algorithm, lightning.LightningModule)
             with trainer.init_module():
+                # A bit hacky, but we have to do this because the lightningmodule isn't associated
+                # with a Trainer.
+                algorithm._device = torch.get_default_device()
                 algorithm.configure_model()
 
         tensor_regression.check(

diff --git a/project/conftest.py b/project/conftest.py
@@ -288,6 +288,9 @@ def algorithm(
     algorithm = instantiate_algorithm(experiment_config.algorithm, datamodule=datamodule)
     if isinstance(trainer, lightning.Trainer) and isinstance(algorithm, lightning.LightningModule):
         with trainer.init_module():
+            # A bit hacky, but we have to do this because the lightningmodule isn't associated
+            # with a Trainer.
+            algorithm._device = torch.get_default_device()
             algorithm.configure_model()
     return algorithm