Load individual elements if state dict load fails

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

@@ -89,7 +89,16 @@ def _load_model(
         policy = cast(TorchPolicy, policy)
 
         for name, mod in modules.items():
-            mod.load_state_dict(saved_state_dict[name])
+            try:
+                missing_keys, _ = mod.load_state_dict(
+                    saved_state_dict[name], strict=False
+                )
+                if missing_keys:
+                    logger.warning(
+                        f"Did not find these keys {missing_keys} in checkpoint. Initializing"
+                    )
+            except (KeyError, TypeError):
+                logger.warning(f"Failed to load for module {name}. Initializing")
 
         if reset_global_steps:
             policy.set_step(0)