move N inside step

nerfstudio-project · Oct 2, 2024 · 5186c47 · 5186c47
1 parent dc19061
commit 5186c47
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/examples/simple_trainer.py b/examples/simple_trainer.py
@@ -764,7 +764,7 @@ def train(self):
             # optimize
             for optimizer in self.optimizers.values():
                 if cfg.visible_adam:
-                    optimizer.step(visibility_mask, gaussian_cnt)
+                    optimizer.step(visibility_mask)
                 else:
                     optimizer.step()
                 optimizer.zero_grad(set_to_none=True)

diff --git a/gsplat/optimizers/selective_adam.py b/gsplat/optimizers/selective_adam.py
@@ -37,15 +37,16 @@ class SelectiveAdam(torch.optim.Adam):
         >>> loss.backward()
 
         >>> # Optimization step with selective updates
-        >>> optimizer.step(visibility=visibility_mask, N=N)
+        >>> optimizer.step(visibility=visibility_mask)
 
     """
 
     def __init__(self, params, eps, betas):
         super().__init__(params=params, eps=eps, betas=betas)
 
     @torch.no_grad()
-    def step(self, visibility, N):
+    def step(self, visibility):
+        N = visibility.numel()
         for group in self.param_groups:
             lr = group["lr"]
             eps = group["eps"]