https://github.com/HealthML/active-segmentation/blob/main/src/datasets/doubly_shuffled_nifti_dataset.py#L499 this line will throw an error when self.dim is not 2.