Address github/pwais comments

nerfstudio-project · Dec 29, 2023 · 73241be · 73241be
1 parent 8f53961
commit 73241be
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py
@@ -209,7 +209,7 @@ def populate_modules(self):
         self.step = 0
 
         self.crop_box: Optional[OrientedBox] = None
-        self.back_color = torch.zeros(3, device="cuda")
+        self.back_color = torch.zeros(3)
 
         self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
             num_cameras=self.num_train_data, device="cpu"
@@ -400,9 +400,9 @@ def refinement_after(self, optimizers: Optimizers, step):
                     for group, param in param_groups.items():
                         self.dup_in_optim(optimizers.optimizers[group], dup_idcs, param, 1)
 
-                    # After a guassian is split into two new gaussians, the original one should be also pruned.
+                    # After a guassian is split into two new gaussians, the original one should also be pruned.
                     splits_mask = torch.cat(
-                        (splits, torch.zeros(nsamps * splits.sum() + dups.sum(), device="cuda", dtype=torch.bool))
+                        (splits, torch.zeros(nsamps * splits.sum() + dups.sum(), device=self.device, dtype=torch.bool))
                     )
                     deleted_mask = self.cull_gaussians(splits_mask)
                     param_groups = self.get_gaussian_param_groups()
@@ -595,7 +595,7 @@ def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
         R = camera.camera_to_worlds[0, :3, :3]  # 3 x 3
         T = camera.camera_to_worlds[0, :3, 3:4]  # 3 x 1
         # flip the z and y axes to align with gsplat conventions
-        R_edit = torch.diag(torch.tensor([1, -1, -1], device="cuda", dtype=R.dtype))
+        R_edit = torch.diag(torch.tensor([1, -1, -1], device=self.device, dtype=R.dtype))
         R = R @ R_edit
         # analytic matrix inverse to get world2camera matrix
         R_inv = R.T

diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py
@@ -536,10 +536,10 @@ def main(self) -> None:
             select = np.logical_and(select, np.isfinite(t).all(axis=1))
             n_after = np.sum(select)
             if n_after < n_before:
-                print(f"{n_before - n_after} NaN/Inf elements in {k}")
+                CONSOLE.print(f"{n_before - n_after} NaN/Inf elements in {k}")
 
         if np.sum(select) < n:
-            print(f"values have NaN/Inf in map_to_tensors, only export {np.sum(select)}/{n}")
+            CONSOLE.print(f"values have NaN/Inf in map_to_tensors, only export {np.sum(select)}/{n}")
             for k, t in map_to_tensors.items():
                 map_to_tensors[k] = map_to_tensors[k][select, :]