fix

amazon-science · Dec 4, 2024 · cc82c83 · cc82c83
1 parent 662415b
commit cc82c83
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -14,5 +14,5 @@ def test_pad_and_stack(tensors: list):
 
     ref = torch.concat(tensors).to(dtype=stacked_and_padded.dtype)
 
-    assert torch.sum(stacked_and_padded) == torch.sum(ref)
+    assert torch.sum(torch.nan_to_num(stacked_and_padded, nan=0)) == torch.sum(ref)
     assert torch.nanmean(stacked_and_padded) == torch.nanmean(ref)