Adding missing named param check on ViT (#5196)

pytorch · Jan 17, 2022 · 09c5ddd · 09c5ddd
1 parent f9bee39
commit 09c5ddd
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torchvision/prototype/models/vision_transformer.py b/torchvision/prototype/models/vision_transformer.py
@@ -11,7 +11,7 @@
 from ...models.vision_transformer import VisionTransformer, interpolate_embeddings  # noqa: F401
 from ._api import WeightsEnum, Weights
 from ._meta import _IMAGENET_CATEGORIES
-from ._utils import handle_legacy_interface
+from ._utils import handle_legacy_interface, _ovewrite_named_param
 
 __all__ = [
     "VisionTransformer",
@@ -111,6 +111,9 @@ def _vision_transformer(
 ) -> VisionTransformer:
     image_size = kwargs.pop("image_size", 224)
 
+    if weights is not None:
+        _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
+
     model = VisionTransformer(
         image_size=image_size,
         patch_size=patch_size,