Keep track of memory allocated by GPUMemoryArena

YOUSIKI · Feb 10, 2022 · 1a6c1ff · 1a6c1ff
1 parent d3729d2
commit 1a6c1ff
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/include/tiny-cuda-nn/gpu_memory.h b/include/tiny-cuda-nn/gpu_memory.h
@@ -456,6 +456,8 @@ class GPUMemoryArena {
 
 	~GPUMemoryArena() {
 		try {
+			total_n_bytes_allocated() -= m_size;
+
 			CUDA_CHECK_THROW(cudaDeviceSynchronize());
 			if (m_base_address) {
 				CU_CHECK_THROW(cuMemUnmap(m_base_address, m_size));
@@ -549,6 +551,8 @@ class GPUMemoryArena {
 		CU_CHECK_THROW(cuMemSetAccess(m_base_address, m_size, &access_desc, 1));
 
 		CUDA_CHECK_THROW(cudaDeviceSynchronize());
+
+		total_n_bytes_allocated() += n_bytes_to_allocate;
 	}
 
 	size_t size() const {