Extend the model_weights dict with sub weights

microsoft · Sep 27, 2024 · 109f94c · 109f94c
1 parent 66abac1
commit 109f94c
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/torchgeo/models/api.py b/torchgeo/models/api.py
@@ -67,6 +67,11 @@
     'vit_small_patch16_224': ViTSmall16_Weights,
 }
 
+for name, weight_enum in _model_weights.items():
+    if isinstance(name, str):
+        for sub_weight_enum in weight_enum:
+            _model_weights[str(sub_weight_enum)] = sub_weight_enum
+
 
 def get_model(name: str, *args: Any, **kwargs: Any) -> nn.Module:
     """Get an instantiated model from its name.