fix cpc (#131)

Lightning-Universe · Jul 29, 2020 · 2fd2969 · 2fd2969
1 parent d794e2a
commit 2fd2969
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pl_bolts/models/self_supervised/cpc/cpc_module.py b/pl_bolts/models/self_supervised/cpc/cpc_module.py
@@ -366,18 +366,21 @@ def add_model_specific_args(parent_parser):
  datamodule = CIFAR10DataModule.from_argparse_args(args)
  datamodule.train_transforms = CPCTrainTransformsCIFAR10()
  datamodule.val_transforms = CPCEvalTransformsCIFAR10()
+ args.patch_size = 8
 
  elif args.dataset == 'stl10':
  datamodule = STL10DataModule.from_argparse_args(args)
  datamodule.train_dataloader = datamodule.train_dataloader_mixed
  datamodule.val_dataloader = datamodule.val_dataloader_mixed
  datamodule.train_transforms = CPCTrainTransformsSTL10()
  datamodule.val_transforms = CPCEvalTransformsSTL10()
+ args.patch_size = 16
 
  elif args.dataset == 'imagenet2012':
  datamodule = SSLImagenetDataModule.from_argparse_args(args)
  datamodule.train_transforms = CPCTrainTransformsImageNet128()
  datamodule.val_transforms = CPCEvalTransformsImageNet128()
+ args.patch_size = 32
 
  model = CPCV2(**vars(args), datamodule=datamodule)
  trainer = pl.Trainer.from_argparse_args(args)