Skip to content

GitLab

Commit 2b2e0b59 authored by Hongkun Yu's avatar Hongkun Yu Committed by A. Unique TensorFlower
Browse files

The tf_data_service override only works for the TaskConfig. 

PiperOrigin-RevId: 339488247
parent 290cd97b
Hide whitespace changes
Inline Side-by-side
Showing with 12 additions and 8 deletions
official/core/train_utils.py
View file @ 2b2e0b59
...@@ -82,15 +82,19 @@ def parse_configuration(flags_obj): ...@@ -82,15 +82,19 @@ def parse_configuration(flags_obj):
'runtime': { 'runtime': {
'tpu': flags_obj.tpu, 'tpu': flags_obj.tpu,
}, },
'task': {
'train_data': {
'tf_data_service_address': flags_obj.tf_data_service,
},
'validation_data': {
'tf_data_service_address': flags_obj.tf_data_service,
}
}
}) })
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,
},
'validation_data': {
'tf_data_service_address': flags_obj.tf_data_service,
}
}
})
# 4. Get the second level of override from `--params_override`. # 4. Get the second level of override from `--params_override`.
# `--params_override` is typically used as a further override over the # `--params_override` is typically used as a further override over the
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment