tests

Lightning-Universe · Sep 17, 2020 · e774c7b · e774c7b
1 parent 09d2f89
commit e774c7b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/pl_bolts/models/self_supervised/swav/swav_finetuner.py b/pl_bolts/models/self_supervised/swav/swav_finetuner.py
@@ -3,7 +3,6 @@
 from argparse import ArgumentParser
 
 from pl_bolts.models.self_supervised.swav.swav_resnet import resnet50, resnet18
-from pl_bolts.datamodules import STL10DataModule, ImagenetDataModule
 
 from pl_bolts.models.self_supervised.ssl_finetuner import SSLFineTuner
 from pl_bolts.models.self_supervised.swav.swav_module import SwAV
@@ -12,6 +11,8 @@
 
 
 def cli_main():  # pragma: no-cover
+    from pl_bolts.datamodules import STL10DataModule, ImagenetDataModule
+
     pl.seed_everything(1234)
 
     parser = ArgumentParser()

diff --git a/pl_bolts/models/self_supervised/swav/swav_module.py b/pl_bolts/models/self_supervised/swav/swav_module.py
@@ -14,7 +14,6 @@
 from torch.optim import Adam, SGD
 
 from pl_bolts.models.self_supervised.swav.swav_resnet import resnet50, resnet18
-from pl_bolts.datamodules import STL10DataModule, ImagenetDataModule
 
 from pl_bolts.transforms.dataset_normalizations import stl10_normalization
 from pl_bolts.models.self_supervised.swav.transforms import SwAVTrainDataTransform, SwAVEvalDataTransform
@@ -472,6 +471,8 @@ def add_model_specific_args(parent_parser):
 
 
 def cli_main():
+    from pl_bolts.datamodules import STL10DataModule, ImagenetDataModule
+
     parser = ArgumentParser()
 
     # model args