Merge pull request #43 from pfxuan/empty-cache

Reduce VRAM consumption
pierotofy · Mar 17, 2024 · 927f6b6 · 927f6b6
2 parents 819fae4 + adb9b25
commit 927f6b6
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/model.cpp b/model.cpp
@@ -5,6 +5,11 @@
 #include "rasterize_gaussians.hpp"
 #include "tensor_math.hpp"
 #include "vendor/gsplat/config.h"
+#ifdef USE_HIP
+#include <c10/hip/HIPCachingAllocator.h>
+#else
+#include <c10/cuda/CUDACachingAllocator.h>
+#endif
 
 torch::Tensor randomQuatTensor(long long n){
     torch::Tensor u = torch::rand(n);
@@ -386,6 +391,11 @@ void Model::afterTrain(int step){
         xysGradNorm = torch::Tensor();
         visCounts = torch::Tensor();
         max2DSize = torch::Tensor();
+#ifdef USE_HIP
+        c10::hip::HIPCachingAllocator::emptyCache();
+#else
+        c10::cuda::CUDACachingAllocator::emptyCache();
+#endif
     }
 }