Commits · c63766d4c72e35b52a4c30a2daf10aa1cc575860 · OpenDAS / TransformerEngine

18 Apr, 2024 1 commit

[JAX] Fixing CI failure due to incorrect use of `static_argnums` in jax.jit (#785) · c63766d4

Alp Dener authored Apr 18, 2024



* fixed static argnums for jax.jit in single gpu encoder test, changed warning filtering for pytest
Signed-off-by: Alp Dener <adener@nvidia.com>

* propagating the fix to the JAX mnist example
Signed-off-by: Alp Dener <adener@nvidia.com>

* fixed missing space ibetween flags i QAA scripts
Signed-off-by: Alp Dener <adener@nvidia.com>

* added TE warnings into the ignore list
Signed-off-by: Alp Dener <adener@nvidia.com>

---------
Signed-off-by: Alp Dener <adener@nvidia.com>

c63766d4