reinstate try-catch for loading missing modules

Signed-off-by: Rohit Jena <[email protected]>
rohitrango · Jul 30, 2024 · 1bd4a73 · 1bd4a73
1 parent 95a8a98
commit 1bd4a73
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/openaimodel.py b/nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/openaimodel.py
@@ -1245,10 +1245,13 @@ def _load_pretrained_model(self, state_dict, ignore_mismatched_sizes=False, from
         ):
             # GroupNormOpt fuses activation function to one layer, thus the indexing of weights are shifted for following
             for key_ in missing_keys:
-                s = key_.split('.')
-                idx = int(s[-2])
-                new_key_ = ".".join(s[:-2] + [str(int(idx + 1))] + [s[-1]])
-                state_dict[key_] = state_dict[new_key_]
+                try:
+                    s = key_.split('.')
+                    idx = int(s[-2])
+                    new_key_ = ".".join(s[:-2] + [str(int(idx + 1))] + [s[-1]])
+                    state_dict[key_] = state_dict[new_key_]
+                except:
+                    continue
 
         loaded_keys = list(state_dict.keys())
         missing_keys = list(set(expected_keys) - set(loaded_keys))