From d491c9e22d3eb5c8424df8e0bfd32204cfa0614d Mon Sep 17 00:00:00 2001 From: jeongukjae Date: Tue, 14 Jun 2022 16:54:46 +0900 Subject: [PATCH] Register addons ops before loading v2 saved model --- tensorflow_transform/saved/saved_transform_io_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow_transform/saved/saved_transform_io_v2.py b/tensorflow_transform/saved/saved_transform_io_v2.py index e046acf0..02bd041d 100644 --- a/tensorflow_transform/saved/saved_transform_io_v2.py +++ b/tensorflow_transform/saved/saved_transform_io_v2.py @@ -121,6 +121,7 @@ def __init__(self, saved_model_dir: str): defined in `../constants.py` ('transform' and 'transform_signature', respectively). """ + saved_transform_io._maybe_register_addon_ops() # pylint: disable=protected-access # TODO(b/160294509): Stop using tf.compat.v2 when TF1.15 support is # dropped. imported = tf.compat.v2.saved_model.load(saved_model_dir) @@ -151,7 +152,6 @@ def __init__(self, saved_model_dir: str): outputs_to_inputs_map = _get_output_to_inputs_map(structured_outputs) self._initialize(load_v2_in_compat, imported, wrapped, structured_inputs, structured_outputs, outputs_to_inputs_map) - saved_transform_io._maybe_register_addon_ops() # pylint: disable=protected-access def _initialize(self, load_v2_in_compat, imported, wrapped, structured_inputs, structured_outputs, outputs_to_inputs_map):