Merge pull request #268 from fmartiescofet/fix_filter_weights

Fix `state_dict` redefinition
IBM · Dec 2, 2024 · 6df109e · 6df109e
2 parents 12e434d + 789a468
commit 6df109e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/terratorch/models/backbones/prithvi_vit.py b/terratorch/models/backbones/prithvi_vit.py
@@ -60,7 +60,7 @@ def checkpoint_filter_fn_vit(
         else:
             clean_dict[k] = v
 
-        state_dict = clean_dict
+    state_dict = clean_dict
 
     state_dict = select_patch_embed_weights(state_dict, model, pretrained_bands, model_bands)
 
@@ -93,7 +93,7 @@ def checkpoint_filter_fn_mae(
         else:
             clean_dict["encoder." + k] = v
 
-        state_dict = clean_dict
+    state_dict = clean_dict
 
     state_dict = select_patch_embed_weights(state_dict, model, pretrained_bands, model_bands)