Try non-strict model loading on RuntimeError. ... (#2214)

* Try non-strict model loading on RuntimeError. ...

This allows tcnn models that contained a torch hash_table to load.
In particular, it allows us to load tcnn models prior to PR #2175.

cf. https://github.com/nerfstudio-project/nerfstudio/pull/2175

* Only allow non-strict loading when called with defaults.

* Reflow to avoid code duplication; blacken.

---------

Co-authored-by: Matthew Tancik <matt@lumalabs.ai>

nerfstudio/pipelines/base_pipeline.py

@@ -106,7 +106,7 @@ class Pipeline(nn.Module):
         """Returns the device that the model is on."""
         return self.model.device
 
-    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
+    def load_state_dict(self, state_dict: Mapping[str, Any], strict: Optional[bool] = None):
         is_ddp_model_state = True
         model_state = {}
         for key, value in state_dict.items():
@@ -122,7 +122,15 @@ class Pipeline(nn.Module):
             model_state = {key[len("module.") :]: value for key, value in model_state.items()}
 
         pipeline_state = {key: value for key, value in state_dict.items() if not key.startswith("_model.")}
-        self.model.load_state_dict(model_state, strict=strict)
+
+        try:
+            self.model.load_state_dict(model_state, strict=True)
+        except RuntimeError:
+            if not strict:
+                self.model.load_state_dict(model_state, strict=False)
+            else:
+                raise
+
         super().load_state_dict(pipeline_state, strict=False)
 
     @profiler.time_function
@@ -410,7 +418,7 @@ class VanillaPipeline(Pipeline):
             (key[len("module.") :] if key.startswith("module.") else key): value for key, value in loaded_state.items()
         }
         self.model.update_to_step(step)
-        self.load_state_dict(state, strict=True)
+        self.load_state_dict(state)
 
     def get_training_callbacks(
         self, training_callback_attributes: TrainingCallbackAttributes