Add epsilon

mpariente · Sep 30, 2024 · 97487fa · 97487fa
1 parent 71b98e8
commit 97487fa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_stoi/stoi.py b/torch_stoi/stoi.py
@@ -215,7 +215,7 @@ def forward(
         # Compute average (E)STOI w. or w/o VAD.
         output = corr_comp.mean(1)
         if self.use_vad:
-            output = output.sum(-1) / mask.sum(-1)
+            output = output.sum(-1) / (mask.sum(-1) + EPS)
         else:
             output = output.mean(-1)