add script to test steps in ai-hub

2026-06-01 16:53:45 -04:00
parent d3ebd2cc5f
commit 090be14a6a
4 changed files with 348 additions and 4 deletions
--- a/examples/training/source/train.py
+++ b/examples/training/source/train.py
@@ -126,10 +126,6 @@ def export_onnx(model: nn.Module, model_dir: Path, image_size: int) -> None:
        do_constant_folding=True,
        input_names=["input"],
        output_names=["logits"],
-        dynamic_axes={
-            "input": {0: "batch_size"},
-            "logits": {0: "batch_size"},
-        },
    )