diff --git a/export/orbax/export/export_manager_obm_test.py b/export/orbax/export/export_manager_obm_test.py index d8e27ae9..0f4c05ae 100644 --- a/export/orbax/export/export_manager_obm_test.py +++ b/export/orbax/export/export_manager_obm_test.py @@ -37,13 +37,6 @@ ) -def _get_expected_tf_function_body(name: str, support_tf_resources: bool): - if support_tf_resources: - return f'inlined_string: "{name}"' - else: - return 'file_system_location { string_path: "%s.pb" }' % name - - class ExportManagerObmTest(parameterized.TestCase, tf.test.TestCase): @parameterized.named_parameters( diff --git a/export/orbax/export/obm_export_test.py b/export/orbax/export/obm_export_test.py index 641f6732..a957707c 100644 --- a/export/orbax/export/obm_export_test.py +++ b/export/orbax/export/obm_export_test.py @@ -38,13 +38,6 @@ def _package_jax_module(m: jax_module.JaxModule): return result -def _get_expected_tf_function_body(name: str, support_tf_resources: bool): - if support_tf_resources: - return f'inlined_string: "{name}"' - else: - return 'file_system_location { string_path: "%s.pb" }' % name - - class ObmExportTest(parameterized.TestCase, tf.test.TestCase): def setUp(self):