From fc4fc5cb15ad994ea82d8c651c9d42172d890de1 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Fri, 18 Oct 2024 11:37:58 -0700 Subject: [PATCH] Properly reset train state after render lock in viewer (#3486) * properly reset train state after render locK * finally block --- nerfstudio/viewer/render_state_machine.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nerfstudio/viewer/render_state_machine.py b/nerfstudio/viewer/render_state_machine.py index cc8002ec9..32448d9bc 100644 --- a/nerfstudio/viewer/render_state_machine.py +++ b/nerfstudio/viewer/render_state_machine.py @@ -148,6 +148,7 @@ class RenderStateMachine(threading.Thread): 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 @@ class RenderStateMachine(threading.Thread): 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):