From b97a42bc328b93c8b1514b107bac81e0dad65abc Mon Sep 17 00:00:00 2001
From: hrouis <houssemeddine.rouis92@gmail.com>
Date: Fri, 17 May 2024 15:40:29 +0200
Subject: [PATCH] fix reshape pybind
---
python_binding/operator/pybind_Reshape.cpp | 2 +-
src/operator/Reshape.cpp | 4 +++-
2 files changed, 4 insertions(+), 2 deletions(-)
diff --git a/python_binding/operator/pybind_Reshape.cpp b/python_binding/operator/pybind_Reshape.cpp
index abda87e62..8d3d4ec07 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 18b35548e..da379fbc9 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) {
--
GitLab