diff --git a/grain/_src/python/dataset/BUILD b/grain/_src/python/dataset/BUILD
index 8fcd0bbcd..718abccc5 100644
--- a/grain/_src/python/dataset/BUILD
+++ b/grain/_src/python/dataset/BUILD
@@ -37,11 +37,12 @@ py_library(
         "dataset.py",
         "//grain/_src/python/dataset/transformations:core_transformations",
     ],
-    srcs_version = "PY3",
     data = select({
         "@platforms//os:windows": ["//grain/_src/python/experimental/index_shuffle/python:index_shuffle_module.pyd"],
         "//conditions:default": ["//grain/_src/python/experimental/index_shuffle/python:index_shuffle_module.so"],
     }),
+    lazy_imports = True,
+    srcs_version = "PY3",
     deps = [
         ":base",
         ":stats",