Properly reset train state after render lock in viewer (#3486)

* properly reset train state after render locK * finally block
nerfstudio-project · Oct 18, 2024 · fc4fc5c · fc4fc5c
1 parent d2f2ad9
commit fc4fc5c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nerfstudio/viewer/render_state_machine.py b/nerfstudio/viewer/render_state_machine.py
@@ -148,6 +148,7 @@ def _render_img(self, camera_state: CameraState):
                         device=self.viewer.get_model().device,
                     )
                     self.viewer.get_model().set_background(background_color)
+                was_training = self.viewer.get_model().training
                 self.viewer.get_model().eval()
                 step = self.viewer.step
                 try:
@@ -168,9 +169,10 @@ def _render_img(self, camera_state: CameraState):
                         with torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt):
                             outputs = self.viewer.get_model().get_outputs_for_camera(camera, obb_box=obb)
                 except viewer_utils.IOChangeException:
-                    self.viewer.get_model().train()
                     raise
-                self.viewer.get_model().train()
+                finally:
+                    if was_training:
+                        self.viewer.get_model().train()
             num_rays = (camera.height * camera.width).item()
             if self.viewer.control_panel.layer_depth:
                 if isinstance(self.viewer.get_model(), SplatfactoModel):