Skip to content

Commit

Permalink
[onert] Fix wrong applied shape in backward of MeanLayer
Browse files Browse the repository at this point in the history
This commit fixes wrong applied shape in backward of MeanLayer.

ONE-DCO-1.0-Signed-off-by: ragmani <ragmani0216@gmail.com>
  • Loading branch information
ragmani committed Jul 4, 2024
1 parent f75846d commit 66bf3de
Showing 1 changed file with 1 addition and 16 deletions.
17 changes: 1 addition & 16 deletions runtime/onert/backend/train/ops/MeanLayer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -48,26 +48,11 @@ void MeanLayer::forward(bool) { cpu::ops::MeanLayer::run(); }

void MeanLayer::backward()
{
nnfw::cker::Shape keep_dim_shape;
if (_keep_dims == false)
{
keep_dim_shape.ReplaceWith(getShape(_input));
auto axes_vec = cpu::ops::getReducerAxes(_axes);
for (const auto &axis : axes_vec)
{
keep_dim_shape.SetDim(axis, 1);
}
}
else
{
keep_dim_shape.ReplaceWith(getShape(_back_prop_input));
}

switch (_back_prop_output->data_type())
{
case OperandType::FLOAT32:
{
nnfw::cker::train::MeanGrad(keep_dim_shape, getBuffer<float>(_back_prop_output),
nnfw::cker::train::MeanGrad(getShape(_back_prop_output), getBuffer<float>(_back_prop_output),
getShape(_back_prop_input), getBuffer<float>(_back_prop_input));
break;
}
Expand Down

0 comments on commit 66bf3de

Please sign in to comment.