diff --git a/src/openpi/training/checkpoints.py b/src/openpi/training/checkpoints.py
index f32a831397..d4273447d8 100644
--- a/src/openpi/training/checkpoints.py
+++ b/src/openpi/training/checkpoints.py
@@ -4,7 +4,7 @@
 import concurrent.futures as futures
 import dataclasses
 import logging
-from typing import Protocol
+from typing import TYPE_CHECKING, Protocol
 
 from etils import epath
 import jax
@@ -13,7 +13,9 @@
 
 from openpi.shared import array_typing as at
 import openpi.shared.normalize as _normalize
-import openpi.training.data_loader as _data_loader
+
+if TYPE_CHECKING:
+    import openpi.training.data_loader as _data_loader
 import openpi.training.utils as training_utils