Added custom_dtype support for wav2vec2

asmigosw · asmigosw · commit cd8cfadee117 · 2026-03-13T12:07:51.000Z
Signed-off-by: asmigosw &lt;asmigosw@qti.qualcomm.com&gt;
Signed-off-by: Dhiraj Kumar Sah &lt;dhirajku@qti.qualcomm.com&gt;
diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py
@@ -4059,7 +4059,7 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:
         seq_len = constants.WAV2VEC2_MAX_SEQ_LEN
 
         example_inputs = {
-            "input_values": torch.zeros((bs, seq_len), dtype=torch.float32),
+            "input_values": torch.zeros((bs, seq_len), dtype=self.model.config.torch_dtype),
         }
 
         dynamic_axes = {"input_values": {0: "batch_size", 1: "seq_len"}}
@@ -4205,6 +4205,8 @@ def cloud_ai_100_feature_generate(
         input_values = np.array(
             torch.nn.functional.pad(inputs["input_values"], (0, self.seq_len - input_ids_len), "constant", 0)
         )
+        needed_dtype = self.model.config.torch_dtype
+        input_values = input_values.astype(DTYPE_TO_STRING_MAP[needed_dtype])
         inputs = dict(input_values=input_values)
         outputs = self.qpc_session.run(inputs)
 

Original file line number	Diff line number	Diff line change
`@@ -4059,7 +4059,7 @@ def export(self, export_dir: Optional[str] = None, **kwargs) -> str:`
`4059`	`4059`	`seq_len = constants.WAV2VEC2_MAX_SEQ_LEN`
`4060`	`4060`
`4061`	`4061`	`example_inputs = {`
`4062`		`- "input_values": torch.zeros((bs, seq_len), dtype=torch.float32),`
	`4062`	`+ "input_values": torch.zeros((bs, seq_len), dtype=self.model.config.torch_dtype),`
`4063`	`4063`	`}`
`4064`	`4064`
`4065`	`4065`	`dynamic_axes = {"input_values": {0: "batch_size", 1: "seq_len"}}`
`@@ -4205,6 +4205,8 @@ def cloud_ai_100_feature_generate(`
`4205`	`4205`	`input_values = np.array(`
`4206`	`4206`	`torch.nn.functional.pad(inputs["input_values"], (0, self.seq_len - input_ids_len), "constant", 0)`
`4207`	`4207`	`)`
	`4208`	`+ needed_dtype = self.model.config.torch_dtype`
	`4209`	`+ input_values = input_values.astype(DTYPE_TO_STRING_MAP[needed_dtype])`
`4208`	`4210`	`inputs = dict(input_values=input_values)`
`4209`	`4211`	`outputs = self.qpc_session.run(inputs)`
`4210`	`4212`