ensure that non_trainable_indices are on the correct device during training

nerfstudio/cameras/camera_optimizers.py

@@ -58,12 +58,14 @@ class CameraOptimizer(nn.Module):
         config: CameraOptimizerConfig,
         num_cameras: int,
         device: Union[torch.device, str],
+        non_trainable_camera_indices: Optional[Int[Tensor, "num_non_trainable_cameras"]] = None,
         **kwargs,
     ) -> None:
         super().__init__()
         self.config = config
         self.num_cameras = num_cameras
         self.device = device
+        self.non_trainable_camera_indices = non_trainable_camera_indices
 
         # Initialize learnable parameters.
         if self.config.mode == "off":
@@ -95,6 +97,11 @@ class CameraOptimizer(nn.Module):
             outputs.append(exp_map_SE3(self.pose_adjustment[indices, :]))
         else:
             assert_never(self.config.mode)
+        # Detach non-trainable indices by setting to identity transform
+        if self.non_trainable_camera_indices is not None:
+            if self.non_trainable_camera_indices.device != self.pose_adjustment.device:
+                self.non_trainable_camera_indices = self.non_trainable_camera_indices.to(self.pose_adjustment.device)
+            outputs[0][self.non_trainable_camera_indices] = torch.eye(4, device=self.pose_adjustment.device)[:3, :4]
 
         # Return: identity if no transforms are needed, otherwise multiply transforms together.
         if len(outputs) == 0:
@@ -114,4 +121,7 @@ class CameraOptimizer(nn.Module):
         """
         Add a regularizer
         """
-        loss_dict['camera_opt_regularizer'] = self.pose_adjustment[:,:3].norm(dim=-1).mean() * self.config.trans_l2_penalty + self.pose_adjustment[:,3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
+        loss_dict["camera_opt_regularizer"] = (
+            self.pose_adjustment[:, :3].norm(dim=-1).mean() * self.config.trans_l2_penalty
+            + self.pose_adjustment[:, 3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
+        )