fix running stage

Lightning-AI · Jan 31, 2021 · e3175b1 · e3175b1
1 parent 238e0f2
commit e3175b1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/models/test_restore.py b/tests/models/test_restore.py
@@ -393,7 +393,7 @@ def assert_good_acc():
         # haven't trained with the new loaded model
         dp_model = new_trainer.model
         dp_model.eval()
-        dp_model.module.running_stage = RunningStage.EVALUATING
+        new_trainer.get_model().running_stage = RunningStage.EVALUATING
 
         dataloader = trainer.train_dataloader
         tpipes.run_prediction(dp_model, dataloader, dp=True)