Free TensorRT engine and context (apache#7702)

neo-ai · May 11, 2021 · 85bad62 · 85bad62
1 parent 119ab67
commit 85bad62
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/runtime/contrib/tensorrt/tensorrt_runtime.cc b/src/runtime/contrib/tensorrt/tensorrt_runtime.cc
@@ -109,6 +109,14 @@ class TensorRTRuntime : public JSONRuntimeBase {
   }
 
 #ifdef TVM_GRAPH_RUNTIME_TENSORRT
+  /*! \brief Destroy engines and contexts. */
+  ~TensorRTRuntime() {
+    for (auto& it : trt_engine_cache_) {
+      it.second.context->destroy();
+      it.second.engine->destroy();
+    }
+  }
+
   /*! \brief Run inference using built engine. */
   void Run() override {
     BuildEngine();