diff --git a/learned_optimization/tasks/parametric/image_mlp_ae.py b/learned_optimization/tasks/parametric/image_mlp_ae.py
index c5026ed..b6494aa 100644
--- a/learned_optimization/tasks/parametric/image_mlp_ae.py
+++ b/learned_optimization/tasks/parametric/image_mlp_ae.py
@@ -95,7 +95,7 @@ def init(self, key: PRNGKey) -> Params:
         image = next(self.datasets.train)["image"]  # pyrefly: ignore[missing-attribute]
         return init_net(key, image)
 
-      def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+      def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
         unused_init_net, apply_net = hk.without_apply_rng(
             hk.transform(_forward))
 
diff --git a/learned_optimization/tasks/parametric/image_mlp_vae.py b/learned_optimization/tasks/parametric/image_mlp_vae.py
index 4fc12b9..bb26b63 100644
--- a/learned_optimization/tasks/parametric/image_mlp_vae.py
+++ b/learned_optimization/tasks/parametric/image_mlp_vae.py
@@ -126,7 +126,7 @@ def init(self, key: PRNGKey) -> Params:
         image = next(self.datasets.train)["image"]  # pyrefly: ignore[missing-attribute]
         return hk.transform(_forward).init(key, image)
 
-      def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+      def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray:  # pyrefly: ignore[bad-override]
         image = data["image"]
         center_fn = lambda img: 2 * (img - 0.5)
         image = jax.lax.cond(task_params["center_data"], center_fn, lambda x: x,
diff --git a/learned_optimization/tasks/parametric/lm_rnn.py b/learned_optimization/tasks/parametric/lm_rnn.py
index 8729001..5937225 100644
--- a/learned_optimization/tasks/parametric/lm_rnn.py
+++ b/learned_optimization/tasks/parametric/lm_rnn.py
@@ -85,7 +85,7 @@ def _forward(inp):
           task_params.values["initializer"])
 
       embed = hk.Embed(vocab_size, self.embed_size, w_init=w_init)(inp)
-      rnn_core = self.rnn_core_fn(self.rnn_size)  # pytype: disable=wrong-arg-count  # trace-all-classes
+      rnn_core = self.rnn_core_fn(self.rnn_size)  # pyrefly: ignore[bad-argument-count]
 
       # Make learnable initial states.
       template_state = rnn_core.initial_state(1)
@@ -124,7 +124,7 @@ def init(self, rng: PRNGKey) -> Params:
         seq = batch["obs"]
         return init_net(rng, seq)
 
-      def loss(self, params: Params, rng: PRNGKey, data: Batch) -> jnp.ndarray:  # pytype: disable=signature-mismatch  # jax-ndarray
+      def loss(self, params: Params, rng: PRNGKey, data: Batch) -> jnp.ndarray:
         net = hk.without_apply_rng(hk.transform(_forward))
         obs = data["obs"]
         target = data["target"]