diff --git a/src/operator/Reshape.cpp b/src/operator/Reshape.cpp index ab53c094dac09879c1bec86509463aab2280ca92..624524b3be49cd8437810d2ed7249f98246365fb 100644 --- a/src/operator/Reshape.cpp +++ b/src/operator/Reshape.cpp @@ -54,6 +54,10 @@ bool Aidge::Reshape_Op::forwardDims(bool /*allowDataDependency*/) { dimSize = 1; negativeIndex = static_cast<DimIdx_t>(i); } + else if (dimSize == 0) + { + dimSize = getInput(0) -> dims()[i]; + } outDims.push_back(static_cast<DimSize_t>(dimSize)); outSize *= static_cast<DimSize_t>(dimSize); }