[neurun] Resolve Division by zero in Softmax layer (#7055)

This commit will resolve division by zero in Softmax layer

Signed-off-by: Seok NamKoong <sk.namkoong@samsung.com>

diff --git a/runtimes/neurun/backend/cpu/kernel/SoftMaxLayer.cc b/runtimes/neurun/backend/cpu/kernel/SoftMaxLayer.cc
index ffeba11..f71a779 100644 (file)
--- a/runtimes/neurun/backend/cpu/kernel/SoftMaxLayer.cc
+++ b/runtimes/neurun/backend/cpu/kernel/SoftMaxLayer.cc
@@ -81,6 +81,9 @@ void SoftMaxLayer::softmaxFloat32()
   if (getNumberOfDimensions(_inputShape) == 2)
   {
     uint32_t batch_size = getSizeOfDimension(_inputShape, 0);
+    if (batch_size == 0)
+      throw std::runtime_error("batch_size should not be 0");
+
     uint32_t input_size = getNumberOfElements(_inputShape) / batch_size;
     Softmax(_inputData.f, input_size, batch_size, _beta, _outputData.f);
   }
@@ -104,6 +107,9 @@ void SoftMaxLayer::softmaxQuant8()
   if (getNumberOfDimensions(_inputShape) == 2)
   {
     uint32_t batch_size = getSizeOfDimension(_inputShape, 0);
+    if (batch_size == 0)
+      throw std::runtime_error("batch_size should not be 0");
+
     uint32_t input_size = getNumberOfElements(_inputShape) / batch_size;
     shapeIn4D.dimensions = {batch_size, 1, 1, input_size};
   }