tests/jax/test_fused_attn.py · 5b6546c89b3a88672f5ac6b3ccdb05f0a6ad5609 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Fix unit tests to work around cuDNN 9.4 regression of 0 length sequences (#1179) · df699655
Michael Goldfarb authored Sep 16, 2024
```
Modify unit tests to work around cuDNN 9.4 regression.
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
```
df699655

test_fused_attn.py 27.4 KB