diff --git a/recml/core/metrics/confusion_metrics.py b/recml/core/metrics/confusion_metrics.py
index 547c24d..7b42e62 100644
--- a/recml/core/metrics/confusion_metrics.py
+++ b/recml/core/metrics/confusion_metrics.py
@@ -594,7 +594,7 @@ def _estimate_confusion_matrix(thresholds: jt.Scalar) -> tuple[
     return jnp.sum(tp), jnp.sum(tn), jnp.sum(fp), jnp.sum(fn)
 
   thresholds = jnp.asarray(thresholds, dtype=jnp.float32)  # pyrefly: ignore[bad-assignment]
-  return jax.vmap(_estimate_confusion_matrix)(thresholds)
+  return jax.vmap(_estimate_confusion_matrix)(thresholds)  # pyrefly: ignore[bad-argument-type]
 
 
 def default_thresholds(num_thresholds: int) -> np.ndarray: