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,
if (val < 0)
val+=static_cast<std::int32_t>(getInput(0)->nbDims());
});
std::sort(axes.cbegin(), axes.cend());
std::sort(axes.begin(), axes.end());
// build output dimensions
std::vector<DimSize_t> outDims = getInput(0)->dims();
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 {
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