examples/jax/encoder/test_model_parallel_encoder.py · 7b18f235465636d54032f676b4c227b2b7f62e3b · OpenDAS / TransformerEngine

Jax example cleanup and replace pjit with jit. (#1107) · 309c6d49

Frédéric Bastien authored Aug 23, 2024



* Use jit instead of pjit

---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

309c6d49

test_model_parallel_encoder.py 16 KB

Replace test_model_parallel_encoder.py