Merge main

pierotofy · Mar 20, 2024 · 8292e30 · 8292e30
2 parents 17211cc + 927f6b6
commit 8292e30
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/model.cpp b/model.cpp
@@ -6,6 +6,12 @@
 #include "tensor_math.hpp"
 #include "gsplat.hpp"
 
+#ifdef USE_HIP
+#include <c10/hip/HIPCachingAllocator.h>
+#elif defined(USE_CUDA)
+#include <c10/cuda/CUDACachingAllocator.h>
+#endif
+
 torch::Tensor randomQuatTensor(long long n){
     torch::Tensor u = torch::rand(n);
     torch::Tensor v = torch::rand(n);
@@ -442,6 +448,11 @@ void Model::afterTrain(int step){
         xysGradNorm = torch::Tensor();
         visCounts = torch::Tensor();
         max2DSize = torch::Tensor();
+#ifdef USE_HIP
+        c10::hip::HIPCachingAllocator::emptyCache();
+#elif defined(USE_CUDA)
+        c10::cuda::CUDACachingAllocator::emptyCache();
+#endif
     }
 }