Free memory on reload using memory allocator clear func (octoml#181)

Lunderberg · May 19, 2023 · 058cbbf · 058cbbf
1 parent 944fc69
commit 058cbbf
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/cpp/llm_chat.cc b/cpp/llm_chat.cc
@@ -477,6 +477,12 @@ class LLMChat {
   }
 
   void Reload(tvm::runtime::Module executable, String model_path) {
+    // Step 0. Clear the previously allocated memory.
+    const PackedFunc* fclear_memory_manager =
+        tvm::runtime::Registry::Get("vm.builtin.memory_manager.clear");
+    ICHECK(fclear_memory_manager) << "Cannot find env function vm.builtin.memory_manager.clear";
+    (*fclear_memory_manager)();
+
     // Step 1. Set tokenizer.
     this->tokenizer_ = TokenizerFromPath(model_path);