Fix bug by changes in Big Vision.

PiperOrigin-RevId: 588003859
google-research · Dec 5, 2023 · dfb9ee0 · dfb9ee0
1 parent 7386d27
commit dfb9ee0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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(