[onert] Implement notifyLastUse in training TensorBuilder

This commit add implementation of the notifyLastUse function in the TensorBuilder for training. ONE-DCO-1.0-Signed-off-by: ragmani <[email protected]>
Samsung · Jul 29, 2024 · 2740e4e · 2740e4e
1 parent 19b0b2b
commit 2740e4e
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/runtime/onert/backend/train/TensorBuilder.cc b/runtime/onert/backend/train/TensorBuilder.cc
@@ -110,9 +110,16 @@ void TensorBuilder::notifyFirstUse(const ir::OperandIndex &index)
   }
 }
 
-void TensorBuilder::notifyLastUse(const ir::OperandIndex &)
+void TensorBuilder::notifyLastUse(const ir::OperandIndex &index)
 {
-  // TODO Support momory plan
+  if (_as_constants[index])
+  {
+    _tensor_mgr->releaseTrainablePlan(index);
+  }
+  else
+  {
+    _tensor_mgr->releaseNonConstPlan(index);
+  }
 }
 
 void TensorBuilder::notifyBackwardFirstUse(const ir::OperandIndex &index)