diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
index d19165082..13733ec1e 100644
--- a/t5x/checkpoints.py
+++ b/t5x/checkpoints.py
@@ -62,6 +62,9 @@
 import tensorstore as ts
 import typing_extensions
 
+from tensorboard.backend.event_processing import directory_watcher
+from tensorboard.backend.event_processing import event_file_loader
+from tensorboard.backend.event_processing import io_wrapper
 
 PartitionSpec = partitioning.PartitionSpec
 PyTree = Any