Fix overflow bug in ByteVocabulary._encode_tf

PiperOrigin-RevId: 570396249
google · Oct 3, 2023 · bae5d7e · bae5d7e
1 parent 29c70c0
commit bae5d7e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/seqio/vocabularies.py b/seqio/vocabularies.py
@@ -643,8 +643,10 @@ def _encode_tf(self, s):
     Returns:
       a 1d tf.Tensor with dtype tf.int32
     """
-    tf_ids = tf.io.decode_raw(s, tf.uint8) + self._num_special_tokens
-    return tf.dtypes.cast(tf_ids, tf.int32)
+    return (
+        tf.dtypes.cast(tf.io.decode_raw(s, tf.uint8), tf.int32)
+        + self._num_special_tokens
+    )
 
   def _decode_tf(self, ids):
     """Decode in TensorFlow.