size

Lightning-Universe · Mar 10, 2021 · abe5c09 · abe5c09
1 parent 0679bab
commit abe5c09
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/pl_bolts/models/self_supervised/byol/byol_module.py b/pl_bolts/models/self_supervised/byol/byol_module.py
@@ -151,7 +151,7 @@ def configure_optimizers(self):
  def add_model_specific_args(parent_parser):
  parser = ArgumentParser(parents=[parent_parser], add_help=False)
  parser.add_argument('--online_ft', action='store_true', help='run online finetuner')
- parser.add_argument('--dataset', type=str, default='cifar10', help='cifar10, imagenet2012, stl10')
+ parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10', 'imagenet2012', 'stl10'])
 
  (args, _) = parser.parse_known_args()
 

diff --git a/pl_bolts/models/self_supervised/moco/moco2_module.py b/pl_bolts/models/self_supervised/moco/moco2_module.py
@@ -320,7 +320,7 @@ def add_model_specific_args(parent_parser):
  parser.add_argument('--momentum', type=float, default=0.9)
  parser.add_argument('--weight_decay', type=float, default=1e-4)
  parser.add_argument('--data_dir', type=str, default='./')
- parser.add_argument('--dataset', type=str, default='cifar10', help='cifar10, stl10, imagenet2012')
+ parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10', 'imagenet2012', 'stl10'])
  parser.add_argument('--batch_size', type=int, default=256)
  parser.add_argument('--use_mlp', action='store_true')
  parser.add_argument('--meta_dir', default='.', type=str, help='path to meta.bin for imagenet')

diff --git a/tests/models/self_supervised/test_ssl_scripts.py b/tests/models/self_supervised/test_ssl_scripts.py
@@ -4,7 +4,7 @@
 
 from tests import _MARK_REQUIRE_GPU, DATASETS_PATH
 
-_DEFAULT_ARGS = f"--data_dir {DATASETS_PATH} --max_epochs 1 --max_steps 4 --batch_size 8 --num_workers 0"
+_DEFAULT_ARGS = f"--data_dir {DATASETS_PATH} --max_epochs 1 --max_steps 2 --batch_size 8 --num_workers 0"
 
 
 # todo: failing for GPU as some is on CPU other on GPU