Skip to content

Commit 2661015

Browse files
authored
Add pattern matching for VIT (#451)
1 parent fddf6b0 commit 2661015

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

src/sparseml/pytorch/utils/quantization/quantize_qat_export.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -745,7 +745,9 @@ def _convert_quantizable_matmul_and_add(model: ModelProto):
745745
continue
746746
if output_quantize_node.op_type != "QuantizeLinear":
747747
continue
748-
bias_initializer = get_init_by_name(model, bias_add_node.input[1])
748+
bias_initializer = get_init_by_name(model, bias_add_node.input[1]) or (
749+
get_init_by_name(model, bias_add_node.input[0])
750+
)
749751
if bias_initializer is None:
750752
continue
751753

0 commit comments

Comments
 (0)