[bug] Fix zero division error

* add edge case handling in cosine_similarity function: in case of zero-valued tensor Signed-off-by: skykongkong8 <[email protected]>
jijoongmoon · Aug 9, 2023 · 24d0d5b · 24d0d5b
1 parent ff5503f
commit 24d0d5b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/nntrainer/layers/fc_layer.cpp b/nntrainer/layers/fc_layer.cpp
@@ -66,7 +66,7 @@ void FullyConnectedLayer::finalize(InitLayerContext &context) {
   context.setEffDimFlagInputDimension(0, 0b1001);
   context.setDynDimFlagInputDimension(0, 0b1000);
 
-  bool is_nchw = (getTensorType() == Tformat::NCHW);
+  bool is_nchw = (context.getFormat() == Tformat::NCHW);
   /** set output dimensions */
   auto const &in_dim = context.getInputDimensions()[0];
   output_dims[0] = in_dim;

diff --git a/test/include/nntrainer_test_util.h b/test/include/nntrainer_test_util.h
@@ -270,6 +270,9 @@ double cosine_similarity(Ta *A /*Predict*/, Tb *B /*Reference */,
     denom_b += ref * ref;
   }
 
+  if (sqrt(denom_a) == 0 && sqrt(denom_b) == 0)
+    return 1;
+
   double cosine_sim = dot / (sqrt(denom_a) * sqrt(denom_b));
   return cosine_sim;
 }
@@ -279,7 +282,7 @@ float mse(Ta *A /* Predicted */, Tb *B /* Reference */, uint32_t size) {
   Ta pred;
   Tb ref;
   float mse_error = 0;
-  for (int i = 0; i < size; i++) {
+  for (uint32_t i = 0; i < size; i++) {
     pred = A[i];
     ref = B[i];
     float diff = pred - ref;