Merge pull request #134 from EthicalML/fix_destroy_staging_tensors

Added destructor for staging tensors
KomputeProject · Feb 7, 2021 · 373ecd0 · 373ecd0
2 parents df5477a + 71f0933
commit 373ecd0
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/OpTensorCreate.cpp b/src/OpTensorCreate.cpp
@@ -23,6 +23,16 @@ OpTensorCreate::OpTensorCreate(
 OpTensorCreate::~OpTensorCreate()
 {
     SPDLOG_DEBUG("Kompute OpTensorCreate destructor started");
+
+    SPDLOG_DEBUG("Kompute OpTensorCreate freeing staging tensors");
+    for (std::shared_ptr<Tensor> tensor : this->mStagingTensors) {
+        if (tensor && tensor->isInit()) {
+            tensor->freeMemoryDestroyGPUResources();
+        } else {
+            SPDLOG_ERROR("Kompute OpTensorCreate expected to free "
+                          "tensor but has already been freed.");
+        }
+    }
 }
 
 void