resolve bug

Author: tchaton

tests/overrides/test_data_parallel.py

@@ -147,6 +147,7 @@ def training_step(self, batch, batch_idx):
     model = TestModel().to(device)
     model.trainer = MagicMock()
     model.trainer._running_stage = RunningStage.TRAINING
+    model.running_stage = RunningStage.TRAINING
     batch = torch.rand(2, 32).to(device)
     batch_idx = 0