ONNX save fix

Signed-off-by: Riyad Islam <[email protected]>

Author: i-riyad

examples/onnx_ptq/download_example_onnx.py

@@ -20,7 +20,7 @@
 import timm
 import torch
 
-from modelopt.torch._deploy.utils import get_onnx_bytes
+from modelopt.torch._deploy.utils import get_onnx_bytes_and_metadata, OnnxBytes
 
 
 def export_to_onnx(model, input_shape, onnx_save_path, device, weights_dtype="fp32"):
@@ -29,15 +29,24 @@ def export_to_onnx(model, input_shape, onnx_save_path, device, weights_dtype="fp
     input_dtype = model.parameters().__next__().dtype
     input_tensor = torch.randn(input_shape, dtype=input_dtype).to(device)
 
-    onnx_model_bytes = get_onnx_bytes(
-        model=model,
+    onnx_bytes, _ = get_onnx_bytes_and_metadata(model=model,
         dummy_input=(input_tensor,),
         weights_dtype=weights_dtype,
     )
-
+    onnx_model = OnnxBytes.from_bytes(onnx_bytes)
+    
     # Write ONNX model to disk
-    with open(onnx_save_path, "wb") as f:
-        f.write(onnx_model_bytes)
+    save_dir = os.path.dirname(os.path.abspath(onnx_save_path))
+    os.makedirs(save_dir, exist_ok=True)
+    
+    for filename, file_bytes in onnx_model.onnx_model.items():
+        if filename.endswith(".onnx"):
+            file_path = onnx_save_path
+        else:
+            file_path = os.path.join(save_dir, filename)
+        with open(file_path, "wb") as f:
+            f.write(file_bytes)
+            print(f"✅ {file_path}")
 
 
 if __name__ == "__main__":