🐛 Fix MNIST test dataloader for shifted data

torch_uncertainty/datamodules/classification/mnist.py

@@ -171,10 +171,12 @@ def test_dataloader(self) -> list[DataLoader]:
 
         Return:
             list[DataLoader]: Dataloaders of the MNIST test set (in
-                distribution data) and FashionMNIST test split
+                distribution data), MNISTC (shifted data), and FashionMNIST test split
                 (out-of-distribution data).
         """
         dataloader = [self._data_loader(self.test)]
         if self.eval_ood:
             dataloader.append(self._data_loader(self.ood))
+        if self.eval_shift:
+            dataloader.append(self._data_loader(self.shift))
         return dataloader

torch_uncertainty/datasets/classification/mnist_c.py

@@ -23,6 +23,7 @@ class MNISTC(VisionDataset):
             takes in the target and transforms it. Defaults to None.
         subset (str): The subset to use, one of ``all`` or the keys in
             ``mnistc_subsets``.
+        shift_severity (int): The shift_severity of the corruption, between 1 and 5.
         download (bool, optional): If True, downloads the dataset from the
             internet and puts it in root directory. If dataset is already
             downloaded, it is not downloaded again. Defaults to False.
@@ -70,6 +71,7 @@ def __init__(
         target_transform: Callable | None = None,
         split: Literal["train", "test"] = "test",
         subset: str = "all",
+        shift_severity: int = 1,
         download: bool = False,
     ) -> None:
         self.root = Path(root)
@@ -90,6 +92,12 @@ def __init__(
             raise ValueError(f"The subset '{subset}' does not exist in MNIST-C.")
         self.subset = subset
 
+        self.shift_severity = shift_severity
+        if shift_severity not in list(range(1, 6)):
+            raise ValueError(
+                "Corruptions shift_severity should be chosen between 1 and 5 " "included."
+            )
+
         if split not in ["train", "test"]:
             raise ValueError(f"The split '{split}' should be either 'train' or 'test'.")
         self.split = split