Free TensorRT engine and context

apache · Mar 19, 2021 · 17dbdb5 · 17dbdb5
1 parent 45442ed
commit 17dbdb5
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();