gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[CI] Moved from assert_allclose to assert_almost_equal (#200)

Browse Source
This commit is contained in:
Philippe Tillet
2021-08-12 12:00:30 -07:00
committed by GitHub
parent 70e28ff380
commit b120d70a0a
5 changed files with 24 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/test/test_cross_entropy.py
View File

@@ -19,7 +19,7 @@ def test_op(M, N, dtype, mode):
tt_y = triton.ops.cross_entropy(x, idx)
th_y = torch.nn.CrossEntropyLoss(reduction="none")(x, idx)
if mode == 'forward':
assert torch.allclose(th_y, tt_y, atol=1e-3, rtol=1e-2)
triton.testing.assert_almost_equal(th_y, tt_y)
# backward pass
elif mode == 'backward':
dy = torch.randn_like(tt_y)
@@ -30,4 +30,4 @@ def test_op(M, N, dtype, mode):
x.grad.zero_()
th_y.backward(dy)
th_dx = x.grad.clone()
assert torch.allclose(th_dx, tt_dx, atol=1e-3, rtol=1e-2)
triton.testing.assert_almost_equal(th_dx, tt_dx)