Fix: fix wrong error message in CelebA datasets when invalid split argument is passed in. (#8866)

Co-authored-by: Nicolas Hug <nh.nicolas.hug@gmail.com>

Author: GdoongMathew

test/test_datasets.py

@@ -532,7 +532,8 @@ def inject_fake_data(self, tmpdir, config):
  self._create_bbox_txt(base_folder, num_images)
  self._create_landmarks_txt(base_folder, num_images)
 
- return dict(num_examples=num_images_per_split[config["split"]], attr_names=attr_names)
+ num_samples = num_images_per_split.get(config["split"], 0) if isinstance(config["split"], str) else 0
+ return dict(num_examples=num_samples, attr_names=attr_names)
 
  def _create_split_txt(self, root):
  num_images_per_split = dict(train=4, valid=3, test=2)
@@ -635,6 +636,28 @@ def test_transforms_v2_wrapper_spawn(self):
  with self.create_dataset(target_type=target_type, transform=v2.Resize(size=expected_size)) as (dataset, _):
  datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
 
+ def test_invalid_split_list(self):
+ with pytest.raises(ValueError, match="Expected type str for argument split, but got type <class 'list'>."):
+ with self.create_dataset(split=[1]):
+ pass
+
+ def test_invalid_split_int(self):
+ with pytest.raises(ValueError, match="Expected type str for argument split, but got type <class 'int'>."):
+ with self.create_dataset(split=1):
+ pass
+
+ def test_invalid_split_value(self):
+ with pytest.raises(
+ ValueError,
+ match="Unknown value '{value}' for argument {arg}. Valid values are {{{valid_values}}}.".format(
+ value="invalid",
+ arg="split",
+ valid_values=("train", "valid", "test", "all"),
+ ),
+ ):
+ with self.create_dataset(split="invalid"):
+ pass
+
 
 class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
  DATASET_CLASS = datasets.VOCSegmentation

torchvision/datasets/celeba.py

@@ -93,7 +93,13 @@ def __init__(
  "test": 2,
  "all": None,
  }
- split_ = split_map[verify_str_arg(split.lower(), "split", ("train", "valid", "test", "all"))]
+ split_ = split_map[
+ verify_str_arg(
+ split.lower() if isinstance(split, str) else split,
+ "split",
+ ("train", "valid", "test", "all"),
+ )
+ ]
  splits = self._load_csv("list_eval_partition.txt")
  identity = self._load_csv("identity_CelebA.txt")
  bbox = self._load_csv("list_bbox_celeba.txt", header=1)