diff options
Diffstat (limited to 'src/backends/neon/NeonLayerSupport.cpp')
-rw-r--r-- | src/backends/neon/NeonLayerSupport.cpp | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/src/backends/neon/NeonLayerSupport.cpp b/src/backends/neon/NeonLayerSupport.cpp index cb709c315a..e856210503 100644 --- a/src/backends/neon/NeonLayerSupport.cpp +++ b/src/backends/neon/NeonLayerSupport.cpp @@ -538,6 +538,12 @@ bool NeonLayerSupport::IsSoftmaxSupported(const TensorInfo& input, const SoftmaxDescriptor& descriptor, Optional<std::string&> reasonIfUnsupported) const { + if (!(descriptor.m_Axis == 1 || + (descriptor.m_Axis < 0 && static_cast<int>(input.GetNumDimensions()) + descriptor.m_Axis == 1))) + { + SetValueChecked(reasonIfUnsupported, "Neon Softmax: Only supports Axis equal to 1."); + return false; + } FORWARD_WORKLOAD_VALIDATE_FUNC(NeonSoftmaxWorkloadValidate, reasonIfUnsupported, input, output, descriptor); } |