aboutsummaryrefslogtreecommitdiff
path: root/delegate/src/Redefine.hpp
diff options
context:
space:
mode:
Diffstat (limited to 'delegate/src/Redefine.hpp')
-rw-r--r--delegate/src/Redefine.hpp7
1 files changed, 5 insertions, 2 deletions
diff --git a/delegate/src/Redefine.hpp b/delegate/src/Redefine.hpp
index fb58ffdf70..91295768d6 100644
--- a/delegate/src/Redefine.hpp
+++ b/delegate/src/Redefine.hpp
@@ -105,10 +105,13 @@ TfLiteStatus VisitReshapeOperator(DelegateData& delegateData,
for (int i=0; i < reshapeOptions->num_dimensions; ++i)
{
targetShape.push_back(reshapeOptions->shape[i]);
- elementCounter = elementCounter * reshapeOptions->shape[i];
+ if (reshapeOptions->shape[i] > 0)
+ {
+ elementCounter = elementCounter * reshapeOptions->shape[i];
+ }
}
// Check the number of elements match, otherwise fall back to using the second input tensor.
- if (elementCounter == inputTensorInfo0.GetNumElements())
+ if (elementCounter <= inputTensorInfo0.GetNumElements())
{
targetShapeFound = true;
}