Add num_samples property to cifar10 dm

pl_bolts/datamodules/base_datamodule.py

@@ -90,18 +90,18 @@ def _split_dataset(self, dataset, train=True):
 
         if train:
             return dataset_train
-        else:
-            return dataset_val
+        return dataset_val
 
     def _get_splits(self, len_dataset):
         if isinstance(self.val_split, int):
             train_len = len_dataset - self.val_split
-            return [train_len, self.val_split]
-
+            splits = [train_len, self.val_split]
         elif isinstance(self.val_split, float):
             val_len = int(self.val_split * len_dataset)
             train_len = len_dataset - val_len
-            return [train_len, val_len]
+            splits = [train_len, val_len]
+
+        return splits
 
     def default_transforms(self):
         return transform_lib.ToTensor()

pl_bolts/datamodules/cifar10_datamodule.py

@@ -96,7 +96,17 @@ def __init__(
             *args,
             **kwargs,
         )
-        self.num_samples = 60000 - val_split
+
+    @property
+    def num_samples(self):
+        len_dataset = 60000
+        if isinstance(self.val_split, int):
+            train_len = len_dataset - self.val_split
+        elif isinstance(self.val_split, float):
+            val_len = int(self.val_split * len_dataset)
+            train_len = len_dataset - val_len
+
+        return train_len
 
     @property
     def num_classes(self):

tests/models/test_classic_ml.py

@@ -41,6 +41,7 @@ def test_logistic_regression_model(tmpdir, datadir):
 
     model = LogisticRegression(input_dim=28 * 28, num_classes=10, learning_rate=0.001)
     model.prepare_data = dm.prepare_data
+    model.setup = dm.setup
     model.train_dataloader = dm.train_dataloader
     model.val_dataloader = dm.val_dataloader
     model.test_dataloader = dm.test_dataloader