[JAX] Fix for Disable FusedAttn with FFI by default (#1304)

rm default value for NVTE_JAX_FUSED_ATTN_WITH_FFI Signed-off-by: Phuong Nguyen <[email protected]>
NVIDIA · Nov 1, 2024 · 4b8ffef · 4b8ffef
1 parent 05c0fb0
commit 4b8ffef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
@@ -379,7 +379,7 @@ def lowering(
 
         wkspace_aval = ctx.avals_out[-1]
 
-        if is_ffi_enabled() and bool(os.getenv("NVTE_JAX_FUSED_ATTN_WITH_FFI", "0")):
+        if is_ffi_enabled() and bool(os.getenv("NVTE_JAX_FUSED_ATTN_WITH_FFI")):
             name = "te_fused_attn_forward_ffi"
             out = ffi.ffi_lowering(name)(
                 ctx,