diff --git a/seqio/evaluation.py b/seqio/evaluation.py
index c770873d..d18b5890 100644
--- a/seqio/evaluation.py
+++ b/seqio/evaluation.py
@@ -33,7 +33,7 @@
 from seqio import metrics as metrics_lib
 from seqio import utils
 import tensorflow.compat.v2 as tf
-import tensorflow_datasets as tfds
+import tensorflow_datasets.public_api as tfds
 import typing_extensions
 
 Task = dataset_providers.Task
diff --git a/seqio/experimental_test.py b/seqio/experimental_test.py
index c1014ed8..e17e55cd 100644
--- a/seqio/experimental_test.py
+++ b/seqio/experimental_test.py
@@ -24,7 +24,7 @@
 from seqio import utils
 from seqio import vocabularies
 import tensorflow.compat.v2 as tf
-import tensorflow_datasets as tfds
+import tensorflow_datasets.public_api as tfds
 
 
 assert_dataset = test_utils.assert_dataset
diff --git a/seqio/loggers.py b/seqio/loggers.py
index 9e659d09..fd6bf8bd 100644
--- a/seqio/loggers.py
+++ b/seqio/loggers.py
@@ -27,7 +27,7 @@
 import numpy as np
 from seqio import metrics as metrics_lib
 import tensorflow as tf
-import tensorflow_datasets as tfds
+import tensorflow_datasets.public_api as tfds