[djl-import] Fixes missing arguments for onnx import (#3458)

deepjavalibrary · Sep 7, 2024 · 61ffdbb · 61ffdbb
1 parent 584a4b1
commit 61ffdbb
Showing 1 changed file with 10 additions and 0 deletions.
diff --git 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",