add assertion

amazon-science · Dec 4, 2024 · 662415b · 662415b
1 parent 357cef4
commit 662415b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -14,4 +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.nanmean(stacked_and_padded) == torch.nanmean(ref)