deepcopy model state_dict in tests (#2887)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lightning-AI · Aug 8, 2020 · 4d0406e · 4d0406e
1 parent 1bb268a
commit 4d0406e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/tests/callbacks/test_early_stopping.py b/tests/callbacks/test_early_stopping.py
@@ -1,3 +1,4 @@
+from copy import deepcopy
 import pickle
 
 import cloudpickle
@@ -24,7 +25,7 @@ def __init__(self, *args, **kwargs):
 
         def on_validation_end(self, trainer, pl_module):
             super().on_validation_end(trainer, pl_module)
-            self.saved_states.append(self.state_dict().copy())
+            self.saved_states.append(deepcopy(self.state_dict()))
 
     class EarlyStoppingTestRestore(EarlyStopping):
         def __init__(self, expected_state):

diff --git a/tests/trainer/test_lr_finder.py b/tests/trainer/test_lr_finder.py
@@ -1,3 +1,4 @@
+from copy import deepcopy
 import pytest
 import torch
 
@@ -33,7 +34,7 @@ def test_model_reset_correctly(tmpdir):
         max_epochs=1,
     )
 
-    before_state_dict = model.state_dict()
+    before_state_dict = deepcopy(model.state_dict())
 
     _ = trainer.lr_find(model, num_training=5)
 

diff --git a/tests/trainer/test_trainer_tricks.py b/tests/trainer/test_trainer_tricks.py
@@ -1,3 +1,4 @@
+from copy import deepcopy
 import pytest
 import torch
 from torch.utils.data import RandomSampler, SequentialSampler, DataLoader
@@ -148,7 +149,7 @@ def test_model_reset_correctly(tmpdir):
         max_epochs=1,
     )
 
-    before_state_dict = model.state_dict()
+    before_state_dict = deepcopy(model.state_dict())
 
     trainer.scale_batch_size(model, max_trials=5)