Merge branch 'main' into analyze_cherry_picks

neuralmagic · Feb 22, 2024 · 509fa1a · 509fa1a
2 parents 08f94c4 + a00ca1e
commit 509fa1a
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/sparsezoo/utils/onnx/external_data.py b/src/sparsezoo/utils/onnx/external_data.py
@@ -300,6 +300,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