From ea0db55765973cacf6984a9eec440ace3d9f2850 Mon Sep 17 00:00:00 2001 From: SeqIO Team Date: Wed, 6 Dec 2023 15:24:58 -0800 Subject: [PATCH] Internal Changes PiperOrigin-RevId: 588560521 --- seqio/dataset_providers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/seqio/dataset_providers.py b/seqio/dataset_providers.py index 3ae616dd..ac9701d2 100644 --- a/seqio/dataset_providers.py +++ b/seqio/dataset_providers.py @@ -1931,6 +1931,7 @@ def get_task_dataset( try_in_mem_cache: bool = True, ) -> tf.data.Dataset: """.""" + def filter_features(ex): return {k: v for k, v in ex.items() if k in output_feature_keys}