fix mnist

Author: Borda

tests/base/datasets.py

@@ -49,13 +49,13 @@ class MNIST(Dataset):
     cache_folder_name = 'complete'
 
     def __init__(self, root: str = PATH_DATASETS, train: bool = True,
-                 normalize: tuple = (0.5, 1.0), download: bool = False):
+                 normalize: tuple = (0.5, 1.0), download: bool = True):
         super().__init__()
         self.root = root
         self.train = train  # training set or test set
         self.normalize = normalize
 
-        self.prepare_data(download)
+        self.prepare_data()
 
         if not self._check_exists(self.cached_folder_path):
             raise RuntimeError('Dataset not found.')
@@ -85,7 +85,7 @@ def _check_exists(self, data_folder: str) -> bool:
             existing = existing and os.path.isfile(os.path.join(data_folder, fname))
         return existing
 
-    def prepare_data(self):
+    def prepare_data(self, download: bool):
         self._download(self.cached_folder_path)
 
     def _download(self, data_folder: str) -> None:
@@ -171,7 +171,7 @@ def _prepare_subset(full_data: torch.Tensor, full_targets: torch.Tensor,
         targets = full_targets[indexes]
         return data, targets
 
-    def prepare_data(self) -> None:
+    def prepare_data(self, download: bool) -> None:
         if self._check_exists(self.cached_folder_path):
             return
         self._download(super().cached_folder_path)