From 8df2ce43344cf69f23670857ccf4e1269fd1ca85 Mon Sep 17 00:00:00 2001 From: Vyacheslav Bazhenov Date: Thu, 19 Dec 2024 12:49:36 +0300 Subject: [PATCH] [onert-micro] Fix MaxPool2D and FullyConnected training kernels This commit fixes errors in MaxPool2D and FullyConnected training kernels. ONE-DCO-1.0-Signed-off-by: Vyacheslav Bazhenov --- onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp | 3 ++- onert-micro/onert-micro/src/train/kernels/MaxPool2D.cpp | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp b/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp index a8248a7fc91..e3035f223bd 100644 --- a/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp +++ b/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp @@ -153,7 +153,8 @@ OMStatus onert_micro::train::train_kernel_CircleFullyConnected(const OMBackpropE // 2. Calculate weight gradient // Init weight grads with zeros - for (int i = 0; i < output_shape.dims(1) * input_shape.dims(1); i += sizeof(float)) + const auto kDlossSizeInBytes = output_shape.dims(1) * input_shape.dims(1) * sizeof(float); + for (int i = 0; i < kDlossSizeInBytes; i += sizeof(float)) *static_cast(static_cast(dloss_dweight_data + i)) = 0; pal::FullyConnectedWeightGrad( diff --git a/onert-micro/onert-micro/src/train/kernels/MaxPool2D.cpp b/onert-micro/onert-micro/src/train/kernels/MaxPool2D.cpp index 096d2d31e8a..8da11dfa3a6 100644 --- a/onert-micro/onert-micro/src/train/kernels/MaxPool2D.cpp +++ b/onert-micro/onert-micro/src/train/kernels/MaxPool2D.cpp @@ -119,7 +119,7 @@ OMStatus onert_micro::train::train_kernel_CircleMaxPool2D(const OMBackpropExecut params.filter_w = options->filter_width(); // Set input grad to zero - for (size_t i = 0; i < input_shape.flatSize(); i += sizeof(float)) + for (size_t i = 0; i < input_shape.flatSize() * sizeof(float); i += sizeof(float)) *static_cast(static_cast(dloss_dinput_data + i)) = 0; // Calculate input grad