Skip to content

Commit

Permalink
relax atol/rtol
Browse files Browse the repository at this point in the history
  • Loading branch information
epwalsh committed Oct 29, 2024
1 parent 0d6597b commit 9609b4b
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions src/test/nn/attention_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,8 +72,13 @@ def test_attention(
y2 = attention(x2)
y = attention(x)

torch.testing.assert_close(y[0:1, :, :], y1)
torch.testing.assert_close(y[1:, :, :], y2)
atol, rtol = None, None
if dtype == torch.bfloat16:
atol = 4e-4
rtol = 5.0

torch.testing.assert_close(y[0:1, :, :], y1, atol=atol, rtol=rtol)
torch.testing.assert_close(y[1:, :, :], y2, atol=atol, rtol=rtol)


@requires_gpu
Expand Down

0 comments on commit 9609b4b

Please sign in to comment.