diff --git a/src/memory_management.c b/src/memory_management.c
index 511b7b8..3bbca9c 100644
--- a/src/memory_management.c
+++ b/src/memory_management.c
@@ -124,7 +124,7 @@ extern "C"
 void gree(void* ptr) {
     #if defined(USE_CUDA) || defined(TEST_MEMORY_MANAGEMENT)
         pthread_mutex_lock(&memory_lock);
-        free_memory(ptr, memory);
+        memory = free_memory(ptr, memory);
         pthread_mutex_unlock(&memory_lock);
     #else
         free(ptr);
diff --git a/src/memory_management.cu b/src/memory_management.cu
index 511b7b8..3bbca9c 100644
--- a/src/memory_management.cu
+++ b/src/memory_management.cu
@@ -124,7 +124,7 @@ extern "C"
 void gree(void* ptr) {
     #if defined(USE_CUDA) || defined(TEST_MEMORY_MANAGEMENT)
         pthread_mutex_lock(&memory_lock);
-        free_memory(ptr, memory);
+        memory = free_memory(ptr, memory);
         pthread_mutex_unlock(&memory_lock);
     #else
         free(ptr);