diff --git a/grain/_src/core/profiler_test.py b/grain/_src/core/profiler_test.py
index 254143d29..b926f42b7 100644
--- a/grain/_src/core/profiler_test.py
+++ b/grain/_src/core/profiler_test.py
@@ -15,7 +15,7 @@
 def _worker_main(worker_init_fn: bytes):
   """Helper function to start a profiler server in a subprocess."""
   worker_init_fn = cloudpickle.loads(worker_init_fn)
-  worker_init_fn()
+  worker_init_fn()  # pyrefly: ignore[not-callable]
   time.sleep(10)
 
 
diff --git a/grain/_src/python/data_loader.py b/grain/_src/python/data_loader.py
index 4bb0f4901..25f584731 100644
--- a/grain/_src/python/data_loader.py
+++ b/grain/_src/python/data_loader.py
@@ -572,8 +572,8 @@ def set_state(self, state: bytes):
     """
     state = json.loads(state.decode())
     if self._validate_state:
-      self._data_loader._validate_state(state)  # pylint: disable=protected-access
-    self._iterator.set_state(state)
+      self._data_loader._validate_state(state)  # pylint: disable=protected-access  # pyrefly: ignore[bad-argument-type]
+    self._iterator.set_state(state)  # pyrefly: ignore[bad-argument-type]
 
   def start_prefetch(self):
     """Starts processing elements asynchronously in the background.