From dfb9ee01ce6dfc5a8228b406e768ee325dd18fcd Mon Sep 17 00:00:00 2001
From: Joan Puigcerver <jpuigcerver@google.com>
Date: Tue, 5 Dec 2023 02:52:38 -0800
Subject: [PATCH] Fix bug by changes in Big Vision.

PiperOrigin-RevId: 588003859
---
 vmoe/data/pp_ops.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vmoe/data/pp_ops.py b/vmoe/data/pp_ops.py
index c8bf68f..809607d 100644
--- a/vmoe/data/pp_ops.py
+++ b/vmoe/data/pp_ops.py
@@ -380,9 +380,9 @@ def tokenize(
     pad_value = tokenizer.string_to_id(pad_value)
 
   def _pp_tokenize(txt):
-    if sample_if_multi:
-      txt = bv_ops_text.ops_general.get_choice(
-          empty_fallback='', key='t')({'t': txt})['t']
+    if sample_if_multi and tf.convert_to_tensor(txt).ndim:
+      txt = bv_ops_text.ops_general.get_choice(key='t')(
+          bv_ops_text.ops_general.get_setdefault('t', '')({'t': txt}))['t']
 
     if lower:
       txt = tf.strings.lower(txt) if sample_if_multi else tf.map_fn(