From 82b3f99f514c8771125b392ca11a5b4bcb45395d Mon Sep 17 00:00:00 2001
From: uyzhang <yi.zhang.4096@gmail.com>
Date: Wed, 24 Apr 2024 18:17:43 +0800
Subject: [PATCH] Fix memory leak in TempAllocator::free()

---
 python/jittor/src/mem/allocator/temp_allocator.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/jittor/src/mem/allocator/temp_allocator.cc b/python/jittor/src/mem/allocator/temp_allocator.cc
index 61213e03..7dae4625 100644
--- a/python/jittor/src/mem/allocator/temp_allocator.cc
+++ b/python/jittor/src/mem/allocator/temp_allocator.cc
@@ -81,6 +81,10 @@ void TempAllocator::free(void* mem_ptr, size_t size, const size_t& allocation) {
         auto it = cached_blocks.lower_bound(get_key(block));
         if (it == cached_blocks.begin()) {
             can_add = false;
+            underlying->free((void*)block->memory_ptr, block->size, 0);
+            unused_memory -= block->size;
+            block_ids.push_back(block->id);
+            delete block;
         } else {
             --it;
             TempCachingBlock* block = it->second;