diff --git a/jax_triton/triton_lib.py b/jax_triton/triton_lib.py
index a98351f..3aaded6 100644
--- a/jax_triton/triton_lib.py
+++ b/jax_triton/triton_lib.py
@@ -703,7 +703,7 @@ def dtype(self) -> np.dtype:
 
 def triton_call(
     *args: jax.Array | bool | int | float | np.float32,
-    kernel: triton.JITFunction,
+    kernel: triton.JITFunction | triton.runtime.Heuristics | triton.runtime.Autotuner,
     out_shape: ShapeDtype | Sequence[ShapeDtype],
     grid: GridOrLambda,
     name: str = "",