Remove deprecated DataModule.dims usage in tests

Author: carmocca

tests/helpers/boring_model.py

@@ -158,18 +158,15 @@ def prepare_data(self):
     def setup(self, stage: Optional[str] = None):
         if stage == "fit" or stage is None:
             self.random_train = Subset(self.random_full, indices=range(64))
-            self.dims = self.random_train[0].shape
 
         if stage in ("fit", "validate") or stage is None:
             self.random_val = Subset(self.random_full, indices=range(64, 64 * 2))
 
         if stage == "test" or stage is None:
             self.random_test = Subset(self.random_full, indices=range(64 * 2, 64 * 3))
-            self.dims = getattr(self, "dims", self.random_test[0].shape)
 
         if stage == "predict" or stage is None:
             self.random_predict = Subset(self.random_full, indices=range(64 * 3, 64 * 4))
-            self.dims = getattr(self, "dims", self.random_predict[0].shape)
 
     def train_dataloader(self):
         return DataLoader(self.random_train)

tests/helpers/datamodules.py

@@ -40,11 +40,6 @@ def __init__(self, data_dir: str = "./", batch_size: int = 32, use_trials: bool
         # TrialMNIST is a constrained MNIST dataset
         self.dataset_cls = TrialMNIST if use_trials else MNIST
 
-        # self.dims is returned when you call dm.size()
-        # Setting default dims here because we know them.
-        # Could optionally be assigned dynamically in dm.setup()
-        self.dims = (1, 28, 28)
-
     def prepare_data(self):
         # download only
         self.dataset_cls(self.data_dir, train=True, download=True)