diff --git a/seqio/dataset_providers_test.py b/seqio/dataset_providers_test.py index 6ef62981..3e1b2cac 100644 --- a/seqio/dataset_providers_test.py +++ b/seqio/dataset_providers_test.py @@ -1932,6 +1932,7 @@ def test_get_dataset_enc_dec_sharded_and_packed(self): assert_dataset(output_ds, expected, expected_dtypes=expected_dtypes) + def register_dummy_task( task_name: str, dataset_fn: Callable[[str, str], tf.data.Dataset],