Small edit

PiperOrigin-RevId: 698144207
google-research · Nov 19, 2024 · 3d18dcd · 3d18dcd
1 parent f02ff83
commit 3d18dcd
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/optformer/common/configurables/factories.py b/optformer/common/configurables/factories.py
@@ -67,7 +67,7 @@ def register_task(
     seqio_dataset_fn: seqio.DatasetFnCallable,
     output_features: dict[str, seqio.Feature],
     distributed: bool = True,
-):
+) -> seqio.Task:
   """Register task from dataset function.
 
   Args:
@@ -76,6 +76,9 @@ def register_task(
       we apply the distributed wrapper, this can be a "dummy".
     output_features: Specifies lengths and preprocessing.
     distributed: Whether the data should come from a reverb server instead.
+
+  Returns:
+    Registered task.
   """
 
   if distributed:
@@ -86,7 +89,7 @@ def register_task(
       splits=(tfds.Split.TRAIN, tfds.Split.VALIDATION, tfds.Split.TEST),
   )
 
-  seqio.TaskRegistry.add(
+  return seqio.TaskRegistry.add(
       task_name,
       source,
       output_features,