From 2dd759f790baf8d183751307cb8dd374c9f9f428 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Wed, 16 Sep 2026 11:16:12 -0700
Subject: [PATCH] suppress pyrefly errors related to JAX scalar types

PiperOrigin-RevId: 982618897
---
 enn/datasets/imagenet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/enn/datasets/imagenet.py b/enn/datasets/imagenet.py
index a8608c0..010b617 100644
--- a/enn/datasets/imagenet.py
+++ b/enn/datasets/imagenet.py
@@ -130,7 +130,7 @@ def load(
     is_training: bool,
     total_batch_size: int,
     data_parallelism: bool,
-    dtype: jnp.dtype = jnp.float32,
+    dtype: jnp.dtype = jnp.float32,  # pyrefly: ignore[bad-function-definition]
     transpose: bool = False,
     fake_data: bool = False,
     image_size: Tuple[int, int] = (224, 224),