Skip to content

Commit eda9435

Browse files
authored
Prevent potential bug in DatasetFolder.make_dataset (#3733)
1 parent ae63bd0 commit eda9435

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

torchvision/datasets/folder.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,13 @@ def make_dataset(
200200
extensions: Optional[Tuple[str, ...]] = None,
201201
is_valid_file: Optional[Callable[[str], bool]] = None,
202202
) -> List[Tuple[str, int]]:
203+
if class_to_idx is None:
204+
# prevent potential bug since make_dataset() would use the class_to_idx logic of the
205+
# find_classes() function, instead of using that of the find_classes() method, which
206+
# is potentially overridden and thus could have a different logic.
207+
raise ValueError(
208+
"The class_to_idx parameter cannot be None."
209+
)
203210
return make_dataset(directory, class_to_idx, extensions=extensions, is_valid_file=is_valid_file)
204211

205212
def find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:

0 commit comments

Comments
 (0)