[fix] Unable to load KV files in FileSystem Saver when using name_sco…

…pe, such as optimizer variable.
tensorflow · Dec 14, 2023 · 10f596c · 10f596c
1 parent 8d68d79
commit 10f596c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py
@@ -174,8 +174,10 @@ def _list_de_variable_saved_files_from_file_system(de_variable_name,
   else:
     keys_pattern = '_mht_*of*_rank*_size{}-keys'.format(proc_size)
     values_pattern = '_mht_*of*_rank*_size{}-values'.format(proc_size)
+  de_variable_saveable_name = string_ops.regex_replace(de_variable_name, "/",
+                                                       "_")
   _shard_name_base_dir = string_ops.string_join(
-      [de_variable_folder_path, de_variable_name], separator='/')
+      [de_variable_folder_path, de_variable_saveable_name], separator='/')
   _shard_name_keys_pattern = string_ops.string_join(
       [_shard_name_base_dir, keys_pattern], separator='')
   shard_keys_file_list = gen_io_ops.matching_files(_shard_name_keys_pattern)