Internal change

PiperOrigin-RevId: 454684592

Internal change
PiperOrigin-RevId: 454684592
99838089 · A. Unique TensorFlower · 969a3f34 · 99838089
Commit 99838089 authored Jun 13, 2022 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/vision/tasks/maskrcnn.py official/vision/tasks/maskrcnn.py +2 -1

No files found.
--- a/official/vision/tasks/maskrcnn.py
+++ b/official/vision/tasks/maskrcnn.py
@@ -122,7 +122,8 @@ class MaskRCNNTask(base_task.Task):
      self,
      params: exp_cfg.DataConfig,
      input_context: Optional[tf.distribute.InputContext] = None,
-      dataset_fn: Optional[dataset_fn_lib.PossibleDatasetType] = None):
+      dataset_fn: Optional[dataset_fn_lib.PossibleDatasetType] = None
+  ) -> tf.data.Dataset:
    """Build input dataset."""
    decoder_cfg = params.decoder.get()
    if params.decoder.type == 'simple_decoder':