diff --git a/grain/_src/python/checkpoint/BUILD b/grain/_src/python/checkpoint/BUILD
index a160c9c46..200572c93 100644
--- a/grain/_src/python/checkpoint/BUILD
+++ b/grain/_src/python/checkpoint/BUILD
@@ -15,6 +15,7 @@ py_library(
 py_library(
     name = "handler",
     srcs = ["handler.py"],
+    lazy_imports = True,
     srcs_version = "PY3",
     deps = [
         "//grain/_src/core:sharding",