This key can be a set

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

terratorch/models/backbones/select_patch_embed_weights.py

@@ -48,16 +48,19 @@ def select_patch_embed_weights(
             _possible_keys_for_proj_weight = {custom_proj_key}
 
         patch_embed_proj_weight_key = state_dict.keys() & _possible_keys_for_proj_weight if (type(state_dict) in [collections.OrderedDict, dict]) else state_dict().keys() & _possible_keys_for_proj_weight
+
         if len(patch_embed_proj_weight_key) == 0:
             msg = "Could not find key for patch embed weight"
             raise Exception(msg)
         if len(patch_embed_proj_weight_key) > 1:
             msg = "Too many matches for key for patch embed weight"
             raise Exception(msg)
-    
+
         # extract the single element from the set
         if isinstance(patch_embed_proj_weight_key, tuple):
             (patch_embed_proj_weight_key,) = patch_embed_proj_weight_key
+        elif isinstance(patch_embed_proj_weight_key, set):
+            patch_embed_proj_weight_key = list(patch_embed_proj_weight_key)[0]
 
         patch_embed_weight = state_dict[patch_embed_proj_weight_key]
 
@@ -80,26 +83,4 @@ def select_patch_embed_weights(
 
         state_dict[patch_embed_proj_weight_key] = temp_weight
 
-    # extract the single element from the set
-    (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()
-
-    # only do this if the patch size and tubelet size match. If not, start with random weights
-    if patch_embed_weights_are_compatible(temp_weight, patch_embed_weight):
-        torch.nn.init.xavier_uniform_(temp_weight.view([temp_weight.shape[0], -1]))
-        for index, band in enumerate(model_bands):
-            if band in pretrained_bands:
-                logging.info(f"Loaded weights for {band} in position {index} of patch embed")
-                temp_weight[:, index] = patch_embed_weight[:, pretrained_bands.index(band)]
-    else:
-        warnings.warn(
-            f"Incompatible shapes between patch embedding of model {temp_weight.shape} and\
-            of checkpoint {patch_embed_weight.shape}",
-            category=UserWarning,
-            stacklevel=1,
-        )
-
-    state_dict[patch_embed_proj_weight_key] = temp_weight
-    return state_dict
+        return state_dict