From 2740e4e102daf8317c07e5ecf832dfcb1e55515c Mon Sep 17 00:00:00 2001 From: ragmani Date: Mon, 29 Jul 2024 17:39:52 +0900 Subject: [PATCH] [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 --- runtime/onert/backend/train/TensorBuilder.cc | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/runtime/onert/backend/train/TensorBuilder.cc b/runtime/onert/backend/train/TensorBuilder.cc index 5c0ac27a7b7..80452858057 100644 --- 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)