fix running stage

Author: awaelchli

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)