Skip to content
Snippets Groups Projects
Commit e941e9fe authored by Maxence Naud's avatar Maxence Naud Committed by Maxence Naud
Browse files

[Fix] set axis conversion to positive values in 'computeOutputDims()'

parent e341bc78
No related branches found
No related tags found
No related merge requests found
...@@ -79,12 +79,12 @@ class ReduceMean_Op : public OperatorTensor, ...@@ -79,12 +79,12 @@ class ReduceMean_Op : public OperatorTensor,
if (val < 0) if (val < 0)
val+=static_cast<std::int32_t>(getInput(0)->nbDims()); val+=static_cast<std::int32_t>(getInput(0)->nbDims());
}); });
std::sort(axes.cbegin(), axes.cend()); std::sort(axes.begin(), axes.end());
// build output dimensions // build output dimensions
std::vector<DimSize_t> outDims = getInput(0)->dims(); std::vector<DimSize_t> outDims = getInput(0)->dims();
if (this->template getAttr<ReduceMeanAttr::KeepDims>()) { if (this->template getAttr<ReduceMeanAttr::KeepDims>()) {
std::for_each(axes.begin(), axes.end(), [&] (const std::int32_t& val) { outDims[val] = 1; }); std::for_each(axes.begin(), axes.end(), [&outDims] (const std::int32_t& val) { outDims[val] = 1; });
} }
else { else {
for (auto it = axes.crbegin(); it != axes.crend(); ++it) for (auto it = axes.crbegin(); it != axes.crend(); ++it)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment