update logic for avoiding div-0; add unit test

Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>

Author: zpqiu

nemo_reinforcer/algorithms/utils.py

@@ -118,13 +118,9 @@ def wrapper(*args, **kwargs):
 
 # need to surpress the masked tensor warnings from pytorch
 @surpress_user_warnings
-def masked_mean(values, mask, dim=None, check_all_zero_mask=True):
+def masked_mean(values, mask, dim=None):
     """Masks values with mask, and computes the mean of the values using the masked values."""
-    if dim is None:
-        if check_all_zero_mask and mask.sum() == 0:
-            return values.sum() * 0
-        return values[mask.bool()].mean()
-    return as_masked_tensor(values, mask.bool()).mean(dim=dim).to_tensor(torch.nan)
+    return (values * mask).sum(dim=dim) / (mask.sum(dim=dim) + 1e-8)
 
 
 def set_seed(seed: int):

tests/unit/algorithms/test_loss_functions.py

@@ -393,10 +393,18 @@ def test_masked_mean_all_zeros():
     values = torch.tensor([1.0, 2.0, 3.0, 4.0])
     mask = torch.zeros_like(values)
 
-    # With check_zero_mask=True (default)
+    # All zeros mask should return 0
     result = masked_mean(values, mask)
-    assert torch.assert_allclose(result, torch.tensor(0.0))
+    print(result)
+    torch.testing.assert_allclose(result, torch.tensor(0.0))
 
     # With check_zero_mask=False
-    result = masked_mean(values, mask, check_all_zero_mask=False)
-    assert torch.isnan(result)  # Should be nan when mask is all zeros
+    mask[0] = 1
+    result = masked_mean(values, mask)
+    torch.testing.assert_allclose(result, torch.tensor(1.0))
+
+    # Case 2: dim is not None
+    values = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    mask = torch.zeros_like(values)
+    result = masked_mean(values, mask, dim=1)
+    torch.testing.assert_allclose(result, torch.tensor([0.0, 0.0]))