Remove feature dimension as dynamic axes for stable diffusion ONNX ex…

…port (#816)
huggingface · Feb 28, 2023 · 536ccac · 536ccac
1 parent bf97d58
commit 536ccac
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -615,8 +615,8 @@ def inputs(self) -> Dict[str, Dict[int, str]]:
     @property
     def outputs(self) -> Dict[str, Dict[int, str]]:
         return {
-            "last_hidden_state": {0: "batch_size", 1: "sequence_length", 2: "feature_dim"},
-            "pooler_output": {0: "batch_size", 1: "feature_dim"},
+            "last_hidden_state": {0: "batch_size", 1: "sequence_length"},
+            "pooler_output": {0: "batch_size"},
         }
 
     def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
@@ -653,7 +653,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]:
         return {
             "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"},
             "timestep": {0: "steps"},
-            "encoder_hidden_states": {0: "batch_size", 1: "sequence_length", 2: "feature_dim"},
+            "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"},
         }
 
     @property