diff --git a/ethosu/regor/tflite/tflite_supported_operators.cpp b/ethosu/regor/tflite/tflite_supported_operators.cpp
index 10643b9f298a36aa1c42bd98ccb18dc77fec0f5d..e5454b26f14cafd0c9d0d202b8dad26cde541506 100644
--- a/ethosu/regor/tflite/tflite_supported_operators.cpp
+++ b/ethosu/regor/tflite/tflite_supported_operators.cpp
@@ -712,7 +712,7 @@ bool TfLiteSupportedOperators::ConstraintSoftmax(const Operation *op)
     auto ifmConn = op->Input(TensorUsage::IFM);
     assert(ifmConn);
     static constexpr int maxProd = 1 << 16;
-    const auto &ifmShape = ifmConn->shape;
+    const auto ifmShape = Shape::PadAxes(ifmConn->shape, 4, 1);
     if ( ifmShape.ElementsWH() > maxProd )
     {
         Failure(op, fmt::format("ifmShape: ({}), W * H = {}", ifmShape.ToString(), ifmShape.ElementsWH()),