Add missing imports & stage

Add missing stage

Author: SeanNaren

pytorch_lightning/trainer/connectors/data_connector.py

@@ -81,7 +81,7 @@ def attach_data(self, model, train_dataloader, val_dataloaders, datamodule):
 
         # set up the passed in dataloaders (if needed)
         self.attach_dataloaders(model, train_dataloader, val_dataloaders)
-        self.attach_datamodule(model, datamodule)
+        self.attach_datamodule(model, datamodule, 'fit')
         self._validate_data_hooks(model)
 
     def __enforce_datamodule_dataloader_override(self, train_dataloader, val_dataloaders, datamodule):

tests/trainer/logging_/test_distributed_logging.py

@@ -11,14 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os
 import platform
 from unittest import mock
+from unittest.mock import Mock
 
 import pytest
 import torch
 
-from pytorch_lightning import Trainer
+from pytorch_lightning import Callback, Trainer
 from tests.helpers import BoringModel