allclose

83b667ef · rusty1s · 20a7cd3c · 83b667ef
Commit 83b667ef authored Mar 10, 2019 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

test/test_std.py test/test_std.py +3 -2

No files found.
--- a/test/test_std.py
+++ b/test/test_std.py
@@ -15,5 +15,6 @@ def test_std(dtype, device, bias):
    index = tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], torch.long, device)
    out = scatter_std(src, index, dim=-1, unbiased=bias)
-    expected = src.std(dim=-1, unbiased=bias)
+    std = src.std(dim=-1, unbiased=bias)[0].item()
-    assert out.tolist() == [[expected[0].item(), 0], [0, expected[0].item()]]
+    expected = torch.tensor([[std, 0], [0, std]], dtype=out.dtype)
+    assert torch.allclose(out, expected)