Fix training when background disabled

zhangkaiitugithub · Nov 13, 2021 · d7a784e · d7a784e
1 parent 71d4730
commit d7a784e
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/svox2/svox2.py b/svox2/svox2.py
@@ -467,6 +467,14 @@ def __init__(
                     dtype=torch.float32, device=device
                 )
             )
+        else:
+            self.background_data = nn.Parameter(
+                torch.empty(
+                    0, 0, 0,
+                    dtype=torch.float32, device=device
+                ),
+                requires_grad=False
+            )
 
         self.register_buffer("links", init_links.view(reso))
         self.links: torch.Tensor
@@ -853,7 +861,7 @@ def volume_render(
                 self.density_data,
                 self.sh_data,
                 basis_data,
-                self.background_data,
+                self.background_data if self.use_background else None,
                 self._to_cpp(replace_basis_data=basis_data),
                 rays._to_cpp(),
                 self.opt._to_cpp(randomize=randomize),