Feat: batch support for onnx and triton compatibility (#239)

* Added batch support Signed-off-by: Siddharth Tyagi <siddhartht@nvidia.com> * file rename Signed-off-by: Siddharth Tyagi <siddhartht@nvidia.com> --------- Signed-off-by: Siddharth Tyagi <siddhartht@nvidia.com> Co-authored-by: Siddharth Tyagi <siddhartht@nvidia.com>
2025-07-26 06:19:52 +05:30
parent d7654ba4e8
commit f1d129d835
2 changed files with 95 additions and 2 deletions
--- a/examples/export.py
+++ b/examples/export.py
@@ -26,8 +26,9 @@ def export_onnx(model, output):
        output_names = [ 'waveform', 'duration' ],
        opset_version = 17, 
        dynamic_axes = {
-            'input_ids': { 1: 'input_ids_len' }, 
-            'waveform': { 0: 'num_samples' }, 
+            'input_ids': {0: "batch_size", 1: 'input_ids_len' }, 
+            'style': {0: "batch_size"}, 
+            "speed": {0: "batch_size"}
        }, 
        do_constant_folding = True, 
    )