Update vdb.py

liruilong940607 · web-flow · commit 8a668783df20 · 2024-06-27T23:23:42.000-07:00
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)
 

Original file line number	Diff line number	Diff line change
`@@ -80,9 +80,8 @@ def __init__(self, init_grid: GridBatch, device="cuda:0") -> None:`
`80`	`80`	`)`
`81`	`81`
`82`	`82`	`def state_dict(self):`
`83`		`- state_dict = super().state_dict()`
	`83`	`+ state_dict = self.state_dict()`
`84`	`84`	`state_dict["grid"] = self.grid`
`85`		`- state_dict["occs"] = self.occs.state_dict()`
`86`	`85`	`return state_dict`
`87`	`86`
`88`	`87`	`def load_state_dict(`
`@@ -96,7 +95,7 @@ def load_state_dict(`
`96`	`95`	`mutable=True,`
`97`	`96`	`)`
`98`	`97`	`remaining_state_dict = {`
`99`		`- k: v for k, v in state_dict.items() if k not in ["grid", "occs"]`
	`98`	`+ k: v for k, v in state_dict.items() if k not in ["grid"]`
`100`	`99`	`}`
`101`	`100`	`super().load_state_dict(remaining_state_dict, strict=strict)`
`102`	`101`