Merge 8a04095c99 into 50e0e3c70c

nerfstudio/engine/trainer.py

@@ -429,7 +429,7 @@ class Trainer:
                 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 @@ class Trainer:
             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"])

nerfstudio/field_components/activations.py

@@ -21,21 +21,21 @@ from typing import TYPE_CHECKING
 import torch
 from jaxtyping import Float
 from torch import Tensor
+from torch.amp import custom_bwd, custom_fwd
 from torch.autograd import Function
-from torch.cuda.amp import custom_bwd, custom_fwd
 
 
 class _TruncExp(Function):
     # Implementation from torch-ngp:
     # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
     @staticmethod
-    @custom_fwd(cast_inputs=torch.float32)
+    @custom_fwd(cast_inputs=torch.float32, device_type="cuda")
     def forward(ctx, x):
         ctx.save_for_backward(x)
         return torch.exp(x)
 
     @staticmethod
-    @custom_bwd
+    @custom_bwd(device_type="cuda")
     def backward(ctx, g):
         x = ctx.saved_tensors[0]
         return g * torch.exp(x.clamp(-15, 15))

nerfstudio/scripts/downloads/download_data.py

@@ -514,7 +514,7 @@ class Mill19Download(DatasetDownload):
             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}")

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