Add missing imports

Lightning-AI · Mar 16, 2021 · 9395139 · 9395139
1 parent dfba137
commit 9395139
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/trainer/logging_/test_distributed_logging.py b/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