diff --git a/nemo_export/onnx_llm_exporter.py b/nemo_export/onnx_llm_exporter.py index 60e9ffbe2..94aab556a 100755 --- a/nemo_export/onnx_llm_exporter.py +++ b/nemo_export/onnx_llm_exporter.py @@ -236,7 +236,7 @@ def _export_to_onnx( f=self.onnx_model_path, input_names=input_names, output_names=output_names, - dynamic_axes={**dynamic_axes_input, **dynamic_axes_output}, + dynamic_axes={**(dynamic_axes_input or {}), **(dynamic_axes_output or {})}, verbose=verbose, opset_version=opset, dynamo=False, diff --git a/tests/unit_tests/export/test_onnx_llm_exporter.py b/tests/unit_tests/export/test_onnx_llm_exporter.py index 4075fc0ee..e87b32544 100644 --- a/tests/unit_tests/export/test_onnx_llm_exporter.py +++ b/tests/unit_tests/export/test_onnx_llm_exporter.py @@ -106,3 +106,23 @@ def test_quantize_without_modelopt(self): pytest.raises(UnavailableError), ): OnnxLLMExporter().quantize(quant_cfg="", forward_loop="") + + def test_export_allows_none_dynamic_axes(self, temp_dir, dummy_model, dummy_tokenizer): + exporter = OnnxLLMExporter( + onnx_model_dir=temp_dir, + model=dummy_model, + tokenizer=dummy_tokenizer, + load_runtime=False, + ) + example_inputs = {"input_ids": torch.randint(0, 10, (2, 4), device=exporter.device)} + with mock.patch("nemo_export.onnx_llm_exporter.torch.onnx.export") as mock_export: + exporter.export( + input_names=["input_ids"], + output_names=["output"], + example_inputs=example_inputs, + dynamic_axes_input=None, + dynamic_axes_output=None, + ) + + _, kwargs = mock_export.call_args + assert kwargs["dynamic_axes"] == {}