Skip to content

Commit

Permalink
When splitting external data, avoid renaming model.data to `model.d…
Browse files Browse the repository at this point in the history
…ata.1` if only one external data file gets eventually saved (#414)
  • Loading branch information
dbogunowicz committed Jan 25, 2024
1 parent e725fa6 commit 1ac94e6
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions src/sparsezoo/utils/onnx/external_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,14 @@ def split_external_data(
# WRITE - ONNX model with updated tensor external data info
onnx.save(model, model_path)

# RENAME - if as a result of splitting we end up with a single file, rename it to
# the original external data file name
if current_external_data_file_number == 1:
os.rename(
os.path.join(base_dir, updated_file_name),
os.path.join(base_dir, external_data_file),
)


def _write_external_data_file_from_base_bytes(
new_file_name, original_byte_ranges, original_file_bytes_reader
Expand Down

0 comments on commit 1ac94e6

Please sign in to comment.