From e2e0f999c8ab22c8625f83c47f0550c41865a8d2 Mon Sep 17 00:00:00 2001
From: Geralt_of_Rivia <87054407+guanjunwu@users.noreply.github.com>
Date: Sun, 3 Mar 2024 15:27:41 +0800
Subject: [PATCH] fix stage rendering error.

---
 gaussian_renderer/__init__.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/gaussian_renderer/__init__.py b/gaussian_renderer/__init__.py
index 80a33da..3c612ee 100644
--- a/gaussian_renderer/__init__.py
+++ b/gaussian_renderer/__init__.py
@@ -77,9 +77,9 @@ def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor,
         scales = pc._scaling
         rotations = pc._rotation
     deformation_point = pc._deformation_table
-    if stage == "coarse" :
+    if "coarse" in stage:
         means3D_final, scales_final, rotations_final, opacity_final, shs_final = means3D, scales, rotations, opacity, shs
-    else:
+    elif "fine" in stage:
         # time0 = get_time()
         # means3D_deform, scales_deform, rotations_deform, opacity_deform = pc._deformation(means3D[deformation_point], scales[deformation_point], 
         #                                                                  rotations[deformation_point], opacity[deformation_point],
@@ -87,9 +87,8 @@ def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor,
         means3D_final, scales_final, rotations_final, opacity_final, shs_final = pc._deformation(means3D, scales, 
                                                                  rotations, opacity, shs,
                                                                  time)
-        # time1 = get_time()
-    # print("deformation forward:",time1-time0)
-    # print(time.max())
+    else:
+        raise NotImplementedError