Fix optimum-cli broke by optimum 1.13.0 release (#217)

fix
huggingface · Sep 8, 2023 · 6fbd1c5 · 6fbd1c5
1 parent 38b3438
commit 6fbd1c5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/optimum/commands/export/neuron.py b/optimum/commands/export/neuron.py
@@ -43,7 +43,7 @@ def parse_args_neuron(parser: "ArgumentParser"):
         default="auto",
         help=(
             "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
-            f" {str(list(TasksManager._TASKS_TO_AUTOMODELS.keys()))}."
+            f" {str(list(TasksManager._TRANSFORMERS_TASKS_TO_MODEL_LOADERS.keys()) + list(TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS.keys()))}."
         ),
     )
     optional_group.add_argument(

diff --git a/optimum/commands/export/neuronx.py b/optimum/commands/export/neuronx.py
@@ -50,7 +50,7 @@ def parse_args_neuronx(parser: "ArgumentParser"):
         default="auto",
         help=(
             "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
-            f" {str(list(TasksManager._TASKS_TO_AUTOMODELS.keys()))}."
+            f" {str(list(TasksManager._TRANSFORMERS_TASKS_TO_MODEL_LOADERS.keys()) + list(TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS.keys()))}."
         ),
     )
     optional_group.add_argument(