This key could not be a tuple

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>
IBM · Jan 7, 2025 · bf0e870 · bf0e870
1 parent 69a60ae
commit bf0e870
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/terratorch/models/backbones/select_patch_embed_weights.py b/terratorch/models/backbones/select_patch_embed_weights.py
@@ -56,7 +56,9 @@ def select_patch_embed_weights(
             raise Exception(msg)
 
         # extract the single element from the set
-        (patch_embed_proj_weight_key,) = patch_embed_proj_weight_key
+        if isinstance(patch_embed_proj_weight_key, tuple):
+            (patch_embed_proj_weight_key,) = patch_embed_proj_weight_key
+
         patch_embed_weight = state_dict[patch_embed_proj_weight_key]
 
         temp_weight = model.state_dict()[patch_embed_proj_weight_key].clone()