diff --git a/mesh_tensorflow/transformer/utils.py b/mesh_tensorflow/transformer/utils.py
index 018a7518..a7f2530a 100644
--- a/mesh_tensorflow/transformer/utils.py
+++ b/mesh_tensorflow/transformer/utils.py
@@ -479,8 +479,8 @@ def _verify_feature_exists(feature_name, should_exist):
 
     # Detokenize in the graph if supported by vocabulary and accelerator.
     def _maybe_detokenize(ids, vocab):
-      if not use_tpu and hasattr(vocab, "decode_tf"):
-        return vocab.decode_tf(ids)
+      if not use_tpu and hasattr(vocab, "tf_tokenizer"):
+        return vocab.tf_tokenizer.detokenize(ids)
       return ids
     if mode == "score":
       # compute log-likelihoods per sequence