ONNXToDT: changes to order of OrderedDict elements

huggingface · Oct 25, 2024 · 69802e9 · 69802e9
1 parent 8387590
commit 69802e9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py
@@ -175,9 +175,9 @@ class OnnxConfig(ExportConfig, ABC):
         ),
         "reinforcement-learning": OrderedDict(
             {
-                "return_preds": {0: "batch_size", 1: "sequence_length"},
-                "action_preds": {0: "batch_size", 1: "sequence_length", 2: "act_dim"},
                 "state_preds": {0: "batch_size", 1: "sequence_length", 2: "state_dim"},
+                "action_preds": {0: "batch_size", 1: "sequence_length", 2: "act_dim"},
+                "return_preds": {0: "batch_size", 1: "sequence_length"},
                 "last_hidden_state": {0: "batch_size", 1: "sequence_length", 2: "last_hidden_state"},
             }
         ),