Empty torch cache after optimizer tensor replacement

Without the cache emptying I got really high VRAM usage and was not able to train beyond 1 million Gaussians with 24GB VRAM. Now I can train to over 5 million Gaussians.
ubc-vision · Dec 22, 2024 · 389bbe4 · 389bbe4
1 parent 491e17a
commit 389bbe4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py
@@ -442,6 +442,8 @@ def replace_tensors_to_optimizer(self, inds=None):
         self._scaling = optimizable_tensors["scaling"]
         self._rotation = optimizable_tensors["rotation"] 
 
+        torch.cuda.empty_cache()
+
         return optimizable_tensors