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

* properly reset train state after render locK

* finally block
This commit is contained in:
Justin Kerr authored and GitHub committed 2024-10-18 11:37:58 -07:00
1 parent d2f2ad9474
commit fc4fc5cb15
1 file changed
+4 -2
+4 -2
View File
@@ -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):