[tensor] Bug fix for average

Bug fix for average to use the correct denominator

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>

diff --git a/nntrainer/src/tensor.cpp b/nntrainer/src/tensor.cpp
index f085755173a80a05453dea43a901608403bb6c6a..94923f6da98589642dcde3be72584a1a8814e318 100644 (file)
--- a/nntrainer/src/tensor.cpp
+++ b/nntrainer/src/tensor.cpp
@@ -887,7 +887,7 @@ Tensor Tensor::average(int axis) const {
 
   Tensor result;
   result = std::move(this->sum(axis));
-  result.divide_i(dim.batch());
+  result.divide_i(dim.getDim()[axis]);
 
   return result;
 }