Merge pull request #372 from IBM/fix-timm-config

Fix timm config loading for prithvi
IBM · Jan 15, 2025 · a34d41f · a34d41f
2 parents a4746d4 + 4d0506d
commit a34d41f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/terratorch/models/backbones/prithvi_vit.py b/terratorch/models/backbones/prithvi_vit.py
@@ -208,7 +208,7 @@ def checkpoint_filter_wrapper_fn(state_dict, model):
                                             f"(pretrained models: {default_cfgs.keys()})")
         # Load pre-trained config from hf
         try:
-            model_args, _ = load_model_config_from_hf(default_cfgs[variant].default.hf_hub_id)
+            model_args = load_model_config_from_hf(default_cfgs[variant].default.hf_hub_id)[0]
             model_args.update(kwargs)
         except:
             logger.warning(f"No pretrained configuration was found on HuggingFace for the model {variant}."