Commit 56453cd5 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change

PiperOrigin-RevId: 471602044
parent 907281c4
...@@ -107,6 +107,12 @@ class ExportConfig(base_config.Config): ...@@ -107,6 +107,12 @@ class ExportConfig(base_config.Config):
def finalize_serving(model_output, export_config): def finalize_serving(model_output, export_config):
"""Adds extra layers based on the provided configuration.""" """Adds extra layers based on the provided configuration."""
if isinstance(model_output, dict):
return {
key: finalize_serving(model_output[key], export_config)
for key in model_output
}
finalize_method = export_config.finalize_method finalize_method = export_config.finalize_method
output_layer = model_output output_layer = model_output
if not finalize_method or finalize_method[0] == 'none': if not finalize_method or finalize_method[0] == 'none':
......
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