fix t5x tests

Signed-off-by: ashors1 <[email protected]>
NVIDIA · Dec 18, 2024 · 034af4f · 034af4f
1 parent 5a74526
commit 034af4f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/.github/container/test-t5x.sh b/.github/container/test-t5x.sh
@@ -175,10 +175,10 @@ seqio.TaskRegistry.add(
     ],
     output_features=dict(
         inputs=seqio.Feature(
-            vocabulary=t5.data.get_default_vocabulary(), add_eos=True, required=False
+            vocabulary=seqio.SentencePieceVocabulary(sentencepiece_model_file="gs://t5-data/vocabs/cc_all.32000.100extra/sentencepiece.model"), add_eos=True, required=False
         ),
         targets=seqio.Feature(
-            vocabulary=t5.data.get_default_vocabulary(), add_eos=True
+            vocabulary=seqio.SentencePieceVocabulary(sentencepiece_model_file="gs://t5-data/vocabs/cc_all.32000.100extra/sentencepiece.model"), add_eos=True
         )
     ),
     metric_fns=[]