From a2f12f354b704fd6f5e2ecae48e7681e3fdb0241 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Sat, 7 Sep 2024 11:12:09 -0700 Subject: [PATCH] [djl-import] Fixes missing arguments for onnx import --- .../src/main/python/djl_converter/arg_parser.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py index a7a0edad382..0b02d4d5ba2 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py +++ b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py @@ -78,6 +78,16 @@ def importer_args(): "--cpu-only", action='store_true', help="Only validate jit traced model on CPU") + parser.add_argument("--optimize", + choices=["O1", "O2", "O3", "O4"], + help="Optimization option for ONNX models") + parser.add_argument( + "--device", + help='The device to use to do the export. Defaults to "cpu".') + parser.add_argument( + "--dtype", + choices=["fp32", "fp16", "bf16"], + help="The floating point precision to use for the export.") group = parser.add_mutually_exclusive_group(required=True) group.add_argument( "-c",