diff --git a/src/onnx/parse_matmulnbits.cpp b/src/onnx/parse_matmulnbits.cpp index af9f09790aa..8356b5d0a57 100644 --- a/src/onnx/parse_matmulnbits.cpp +++ b/src/onnx/parse_matmulnbits.cpp @@ -67,7 +67,15 @@ struct parse_matmulnbits : op_parser ". Actual dims: " + to_string_range(args[1]->get_shape().lens())); std::vector expected_scales_lens{n * n_blocks_per_col}; - if(args[2]->get_shape().lens() != expected_scales_lens) + + // Reshape anything larger than 1 dimension into a 1d tensor so we can check if we have the right amount of elements. + auto scale_input = args[2]; + if(scale_input->get_shape().lens().size() > 1) + { + scale_input = info.add_instruction(make_op("reshape", {{"dims", {scale_input->get_shape().elements()}}}), scale_input); + } + + if(scale_input->get_shape().lens() != expected_scales_lens) MIGRAPHX_THROW("MatMulNBits: Input scales does not match expected dims: " + to_string_range(expected_scales_lens) + ". Actual dims: " + to_string_range(args[2]->get_shape().lens()));