transformer_engine/jax/fused_attn.py · c898ab1bb639cee1773ff7ca5a152c2d63afec8a · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Add checkpoint_name for the recompute granularity control (#542) · c898ab1b
zlsh80826 authored Dec 05, 2023
```
Add checkpoint_name
Signed-off-by: Reese Wang <rewang@nvidia.com>
```
c898ab1b

fused_attn.py 9.48 KB