diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 363c8c9ea..11ba841e0 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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