Allow kwargs in all generate_dummy_inputs() methods (#638)

fix generate_dummy_inputs args
huggingface · Dec 23, 2022 · 433f477 · 433f477
1 parent 874d2b2
commit 433f477
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -555,8 +555,8 @@ def outputs(self) -> Mapping[str, Mapping[int, str]]:
  "pooler_output": {0: "batch_size", 1: "feature_dim"},
  }
 
- def generate_dummy_inputs(self, framework: str = "pt"):
- dummy_inputs = super().generate_dummy_inputs(framework=framework)
+ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
+ dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs)
  if framework == "pt":
  import torch
 
@@ -599,8 +599,8 @@ def outputs(self) -> Mapping[str, Mapping[int, str]]:
  def output_names_for_validation(self, reference_output_names: List[str]) -> List[str]:
  return ["sample"]
 
- def generate_dummy_inputs(self, framework: str = "pt"):
- dummy_inputs = super().generate_dummy_inputs(framework=framework)
+ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
+ dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs)
  dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0]
  return dummy_inputs