[test] Fix bfloat16 tolerance in test_selective_state_update (#20345)

2026-03-16 09:32:01 +05:30
parent 97b2a89334
commit 223a4e485b
1 changed files with 1 additions and 1 deletions
--- a/test/registered/layers/mamba/test_mamba_ssm.py
+++ b/test/registered/layers/mamba/test_mamba_ssm.py
@@ -99,7 +99,7 @@ def test_selective_state_update(dim, dstate, has_z, itype):

    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
    if itype == torch.bfloat16:
-        rtol, atol = 1e-2, 5e-2
+        rtol, atol = 1e-2, 1e-1
        if torch.version.hip:
            atol *= 2
    # set seed