diff --git a/vmoe/checkpoints/base.py b/vmoe/checkpoints/base.py
index 0ceb814..7877c90 100644
--- a/vmoe/checkpoints/base.py
+++ b/vmoe/checkpoints/base.py
@@ -118,7 +118,7 @@ def _parse_step_and_suffix_or_error(filepath):
       sorted(steps_and_suffixes - incomplete_steps_and_suffixes,
              reverse=decreasing),
       lambda x: x[0]):
-    if not suffixes.difference(set(x[1] for x in group)):  # pyrefly: ignore[missing-attribute]
+    if not suffixes.difference(set(x[1] for x in group)):
       yield step
 
 
diff --git a/vmoe/train/optimizer.py b/vmoe/train/optimizer.py
index abb688e..431d461 100644
--- a/vmoe/train/optimizer.py
+++ b/vmoe/train/optimizer.py
@@ -166,7 +166,7 @@ def frozen_fn(params: optax.Params) -> PyTree:
     flatparams = flax.traverse_util.flatten_dict(
         flax.serialization.to_state_dict(params), sep='/')
     output = {
-        key: search_true if pattern.search(key) else search_false  # pyrefly: ignore[missing-attribute]
+        key: search_true if pattern.search(key) else search_false
         for key, value in flatparams.items()
     }
     return flax.serialization.from_state_dict(