From e774c7bb7f89f7b14cc0a8c18bd0dab3665f3555 Mon Sep 17 00:00:00 2001 From: ananyahjha93 Date: Wed, 16 Sep 2020 18:42:09 -0400 Subject: [PATCH] tests --- pl_bolts/models/self_supervised/swav/swav_finetuner.py | 3 ++- pl_bolts/models/self_supervised/swav/swav_module.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pl_bolts/models/self_supervised/swav/swav_finetuner.py b/pl_bolts/models/self_supervised/swav/swav_finetuner.py index d7b673352c..84b708dfb2 100644 --- 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 index 0e77cfe5a6..5a95869515 100644 --- 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