Refactor simclr

Lightning-Universe · Nov 22, 2020 · 64cf526 · 64cf526
1 parent e7966ac
commit 64cf526
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pl_bolts/models/self_supervised/simclr/simclr_module.py b/pl_bolts/models/self_supervised/simclr/simclr_module.py
@@ -331,7 +331,7 @@ def add_model_specific_args(parent_parser):
         parser.add_argument("--gaussian_blur", action="store_true", help="add gaussian blur")
         parser.add_argument("--jitter_strength", type=float, default=1.0, help="jitter strength")
         parser.add_argument("--dataset", type=str, default="cifar10", help="stl10, cifar10")
-        parser.add_argument("--data_path", type=str, default=".", help="path to download data")
+        parser.add_argument("--data_dir", type=str, default=".", help="path to download data")
 
         # training params
         parser.add_argument("--fast_dev_run", action='store_true')
@@ -368,7 +368,7 @@ def cli_main():
 
     if args.dataset == 'stl10':
         dm = STL10DataModule(
-            data_dir=args.data_path,
+            data_dir=args.data_dir,
             batch_size=args.batch_size,
             num_workers=args.num_workers
         )
@@ -391,7 +391,7 @@ def cli_main():
             val_split = args.nodes * args.gpus * args.batch_size
 
         dm = CIFAR10DataModule(
-            data_dir=args.data_path,
+            data_dir=args.data_dir,
             batch_size=args.batch_size,
             num_workers=args.num_workers,
             val_split=val_split
@@ -429,7 +429,7 @@ def cli_main():
         args.online_ft = True
 
         dm = ImagenetDataModule(
-            data_dir=args.data_path,
+            data_dir=args.data_dir,
             batch_size=args.batch_size,
             num_workers=args.num_workers
         )