Fix bug related to refine_scale2d

nerfstudio-project · Aug 21, 2024 · b9a1c37 · b9a1c37
1 parent 435305e
commit b9a1c37
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/gsplat/strategy/default.py b/gsplat/strategy/default.py
@@ -175,6 +175,7 @@ def step_post_backward(
             # reset running stats
             state["grad2d"].zero_()
             state["count"].zero_()
+            state["radii"].zero_()
             torch.cuda.empty_cache()
 
         if step % self.reset_every == 0:
@@ -258,9 +259,9 @@ def _grow_gs(
         n_dupli = is_dupli.sum().item()
 
         is_large = ~is_small
-        if step < self.refine_scale2d_stop_iter:
-            is_large |= state["radii"] > self.grow_scale2d
         is_split = is_grad_high & is_large
+        if step < self.refine_scale2d_stop_iter:
+            is_split |= state["radii"] > self.grow_scale2d
         n_split = is_split.sum().item()
 
         # first duplicate