remove default value from LabelToOneHot

torchvision/prototype/transforms/_type_conversion.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
@@ -15,14 +16,32 @@
 class LabelToOneHot(Transform):
     _transformed_types = (datapoints.Label,)
 
-    def __init__(self, num_categories: int = -1):
+    def __init__(self, num_categories: Optional[int] = None):
         super().__init__()
+        if not ((isinstance(num_categories, int) and num_categories > 0) or num_categories is None):
+            raise ValueError(
+                f"`num_categories` can either be a positive integer or `None`, but got {num_categories} instead."
+            )
         self.num_categories = num_categories
 
     def _transform(self, inpt: datapoints.Label, params: Dict[str, Any]) -> datapoints.OneHotLabel:
-        num_categories = self.num_categories
-        if num_categories == -1 and inpt.categories is not None:
+        if self.num_categories is None and inpt.categories is None:
+            raise RuntimeError(
+                "Can't determine the number of categories, "
+                "since neither `num_categories` on this transform, nor the `.categories` attribute on the label is set!"
+            )
+        elif inpt.categories is None:
+            num_categories = self.num_categories
+        elif self.num_categories is None:
             num_categories = len(inpt.categories)
+        else:
+            num_categories = self.num_categories
+            if num_categories != len(inpt.categories):
+                warnings.warn(
+                    f"`num_categories` set on this transform mismatches the `.categories` attribute on the label: "
+                    f"{num_categories} != {len(inpt.categories)}"
+                )
+
         output = one_hot(inpt.as_subclass(torch.Tensor), num_classes=num_categories)
         return datapoints.OneHotLabel(output, categories=inpt.categories)