Update vdb.py

nerfstudio-project · Jun 28, 2024 · 8a66878 · 8a66878
1 parent 1e46cb1
commit 8a66878
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/nerfacc/estimators/vdb.py b/nerfacc/estimators/vdb.py
@@ -80,9 +80,8 @@ def __init__(self, init_grid: GridBatch, device="cuda:0") -> None:
         )
 
     def state_dict(self):
-        state_dict = super().state_dict()
+        state_dict = self.state_dict()
         state_dict["grid"] = self.grid
-        state_dict["occs"] = self.occs.state_dict()
         return state_dict
 
     def load_state_dict(
@@ -96,7 +95,7 @@ def load_state_dict(
             mutable=True,
         )
         remaining_state_dict = {
-            k: v for k, v in state_dict.items() if k not in ["grid", "occs"]
+            k: v for k, v in state_dict.items() if k not in ["grid"]
         }
         super().load_state_dict(remaining_state_dict, strict=strict)