From 90f53c3d914d4e814a963796053b097201446412 Mon Sep 17 00:00:00 2001
From: Gallasko <gallasko@gmail.com>
Date: Thu, 3 Apr 2025 16:32:19 +0200
Subject: [PATCH] fix: Only allow possitive value of variances in batchnorm
---
aidge_export_cpp/kernels/batchnorm.hpp | 9 ++++++++-
1 file changed, 8 insertions(+), 1 deletion(-)
diff --git a/aidge_export_cpp/kernels/batchnorm.hpp b/aidge_export_cpp/kernels/batchnorm.hpp
index 0260d42..01104f9 100644
--- a/aidge_export_cpp/kernels/batchnorm.hpp
+++ b/aidge_export_cpp/kernels/batchnorm.hpp
@@ -26,7 +26,14 @@ void batchnorm_forward (
const Rescaling_T& __restrict rescaling)
{
for (unsigned int output = 0; output < NB_OUTPUTS; ++output) {
- const Output_T var = sqrt(variances[output] + epsilon);
+ // If the variance is 0, we need to avoid division by 0
+ const Output_T var = epsilon;
+
+ // If the variance is negative, we need to set it to 0 to avoid a sqrt of a negative number
+ if (variances[output] > 0.0)
+ {
+ var = sqrt(variances[output] + epsilon);
+ }
for (int oy = 0; oy < OUTPUTS_HEIGHT; ++oy) {
for (int ox = 0; ox < OUTPUTS_WIDTH; ++ox) {
--
GitLab