Load individual elements if state dict load fails

ml-agents/mlagents/trainers/model_saver/torch_model_saver.py

@@ -89,7 +89,25 @@ def _load_model(
         policy = cast(TorchPolicy, policy)
 
         for name, mod in modules.items():
-            mod.load_state_dict(saved_state_dict[name])
+            try:
+                mod.load_state_dict(saved_state_dict[name])
+            except Exception:
+                if name in saved_state_dict:
+                    logger.warning(f"Failed to load directly for module {name}.")
+                    for mod_element in mod.state_dict():
+                        try:
+                            logger.warning(f"Copying {mod_element}")
+                            mod.state_dict()[mod_element].copy_(
+                                saved_state_dict[name][mod_element]
+                            )
+                        except Exception:
+                            logger.warning(
+                                f"{mod_element} was not found or is not loadable (changed shape). Initializing."
+                            )
+                else:
+                    logger.warning(
+                        f"The module {name} was not found in the checkpoint. Initializing."
+                    )
 
         if reset_global_steps:
             policy.set_step(0)