set running stage on module

Author: awaelchli

tests/overrides/test_data_parallel.py

@@ -156,7 +156,7 @@ def training_step(self, batch, batch_idx):
     pl_module = DeviceAccessModel()
     # required for redirecting the forward call to training_step
     pl_module.trainer = Mock()
-    pl_module.trainer._running_stage = RunningStage.TRAINING
+    pl_module.running_stage = RunningStage.TRAINING
 
     root_device = torch.device("cuda", 0)
     wrapped_module = LightningParallelModule(pl_module).to(root_device)