The tf_data_service override only works for the TaskConfig.

PiperOrigin-RevId: 339488247

The tf_data_service override only works for the TaskConfig.
PiperOrigin-RevId: 339488247
606e3331 · Hongkun Yu · A. Unique TensorFlower · bddc4930 · 606e3331
Commit 606e3331 authored Oct 28, 2020 by Hongkun Yu Committed by A. Unique TensorFlower Oct 28, 2020
Show whitespace changes
Inline Side-by-side

Showing with 12 additions and 8 deletions

official/core/train_utils.py official/core/train_utils.py +12 -8

No files found.
--- a/official/core/train_utils.py
+++ b/official/core/train_utils.py
@@ -82,6 +82,10 @@ def parse_configuration(flags_obj):
      'runtime': {
          'tpu': flags_obj.tpu,
      },
+  })
+  if flags_obj.tf_data_service and isinstance(params.task,
+                                              config_definitions.TaskConfig):
+    params.override({
        'task': {
            'train_data': {
                'tf_data_service_address': flags_obj.tf_data_service,