Set precision as an static arg

Signed-off-by: Ming Huang <[email protected]>
NVIDIA · Dec 15, 2023 · 2ce5724 · 2ce5724
1 parent ca3f0a5
commit 2ce5724
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/jax/dot.py b/transformer_engine/jax/dot.py
@@ -57,7 +57,7 @@ def dequantize(x, dq_dtype, scale_inv):
 
 
 # Apply jit to guarantee correctness of FP8 GEMM.
-@partial(jax.jit, static_argnums=(4, 5))
+@partial(jax.jit, static_argnums=(4, 5, 6))
 def fp8_dot_impl(
         q_lhs: jnp.ndarray,
         q_rhs: jnp.ndarray,