diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a653e1de8d..7f2f0d34d0 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -429,7 +429,7 @@ def _load_checkpoint(self) -> None: load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1] load_path: Path = load_dir / f"step-{load_step:09d}.ckpt" assert load_path.exists(), f"Checkpoint {load_path} does not exist" - loaded_state = torch.load(load_path, map_location="cpu") + loaded_state = torch.load(load_path, map_location="cpu", weights_only=False) self._start_step = loaded_state["step"] + 1 # load the checkpoints for pipeline, optimizers, and gradient scalar self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) @@ -440,7 +440,7 @@ def _load_checkpoint(self) -> None: CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_path}") elif load_checkpoint is not None: assert load_checkpoint.exists(), f"Checkpoint {load_checkpoint} does not exist" - loaded_state = torch.load(load_checkpoint, map_location="cpu") + loaded_state = torch.load(load_checkpoint, map_location="cpu", weights_only=False) self._start_step = loaded_state["step"] + 1 # load the checkpoints for pipeline, optimizers, and gradient scalar self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) diff --git a/nerfstudio/scripts/downloads/download_data.py b/nerfstudio/scripts/downloads/download_data.py index a2c9d27584..653fbf94d0 100644 --- a/nerfstudio/scripts/downloads/download_data.py +++ b/nerfstudio/scripts/downloads/download_data.py @@ -514,7 +514,7 @@ def download(self, save_dir: Path) -> None: split_filepaths = [] for image_path, new_image_path in copied_images.items(): metadata_path = image_path.parent.parent / "metadata" / f"{image_path.stem}.pt" - metadata = torch.load(metadata_path, map_location="cpu") + metadata = torch.load(metadata_path, map_location="cpu", weights_only=False) c2w = torch.eye(4) c2w[:3] = metadata["c2w"] file_path = str(Path("images") / f"{new_image_path.name}") diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 11a8b23416..8d07fe842a 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -59,7 +59,7 @@ def eval_load_checkpoint(config: TrainerConfig, pipeline: Pipeline) -> Tuple[Pat load_step = config.load_step load_path = config.load_dir / f"step-{load_step:09d}.ckpt" assert load_path.exists(), f"Checkpoint {load_path} does not exist" - loaded_state = torch.load(load_path, map_location="cpu") + loaded_state = torch.load(load_path, map_location="cpu", weights_only=False) pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) CONSOLE.print(f":white_check_mark: Done loading checkpoint from {load_path}") return load_path, load_step