diff --git a/python_binding/operator/pybind_Reshape.cpp b/python_binding/operator/pybind_Reshape.cpp index abda87e6249fb783830eacd3ac655299128a42ae..8d3d4ec07229d62bc7908793f845f560e43033b8 100644 --- a/python_binding/operator/pybind_Reshape.cpp +++ b/python_binding/operator/pybind_Reshape.cpp @@ -23,6 +23,6 @@ void init_Reshape(py::module& m) { .def("get_inputs_name", &Reshape_Op::getInputsName) .def("get_outputs_name", &Reshape_Op::getOutputsName); declare_registrable<Reshape_Op>(m, "ReshapeOp"); - m.def("Reshape", &Reshape, py::arg("shape") = std::vector<std::int64_t>(), py::arg("name") = ""); + m.def("Reshape", &Reshape, py::arg("shape") = std::vector<std::int64_t>(), py::arg("allowzero") = false, py::arg("name") = ""); } } // namespace Aidge diff --git a/src/operator/Reshape.cpp b/src/operator/Reshape.cpp index 18b35548e82bdb0329724f828e4d27d8dd178c44..da379fbc9fa364dde475900bc40a82a5bdec19fb 100644 --- a/src/operator/Reshape.cpp +++ b/src/operator/Reshape.cpp @@ -97,7 +97,9 @@ bool Aidge::Reshape_Op::forwardDims(bool /*allowDataDependency*/) { dimSize = getInput(0) -> dims()[i]; } outDims.push_back(static_cast<DimSize_t>(dimSize)); - outSize *= static_cast<DimSize_t>(dimSize); + if (dimSize != 0) { + outSize *= static_cast<DimSize_t>(dimSize); + } } if (foundNegativeDimension) {