diff --git a/seqio/vocabularies.py b/seqio/vocabularies.py
index bf9256ac..5061371a 100644
--- a/seqio/vocabularies.py
+++ b/seqio/vocabularies.py
@@ -22,7 +22,7 @@
 from typing import Any, ClassVar, Dict, Iterable, Optional, Sequence, Union
 
 from absl import logging
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_text as tf_text
 
 from sentencepiece import sentencepiece_model_pb2