diff options
| -rw-r--r-- | src/armnnTfLiteParser/TfLiteParser.cpp | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/src/armnnTfLiteParser/TfLiteParser.cpp b/src/armnnTfLiteParser/TfLiteParser.cpp index 66dc13ef49..9e8af66b49 100644 --- a/src/armnnTfLiteParser/TfLiteParser.cpp +++ b/src/armnnTfLiteParser/TfLiteParser.cpp @@ -2255,6 +2255,9 @@ void TfLiteParserImpl::ParseSqueeze(size_t subgraphIndex, size_t operatorIndex) ReshapeDescriptor reshapeDesc; reshapeDesc.m_TargetShape = outputTensorInfo.GetShape(); + auto outputTensorIds = GetOutputTensorIds(m_Model, subgraphIndex, operatorIndex); + m_TensorInfos[outputTensorIds[0]] = outputTensorInfo; + IConnectableLayer* layer = m_Network->AddReshapeLayer(reshapeDesc, layerName.c_str()); ARMNN_ASSERT(layer != nullptr); layer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo); |