MatMulNBits collapse shape input when > 1d (#3698)

ROCm · Dec 11, 2024 · 7eaafc3 · 7eaafc3
1 parent 64fe0c5
commit 7eaafc3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/onnx/parse_matmulnbits.cpp b/src/onnx/parse_matmulnbits.cpp
@@ -66,10 +66,10 @@ struct parse_matmulnbits : op_parser<parse_matmulnbits>
                            to_string_range(expected_b_lens) +
                            ". Actual dims: " + to_string_range(args[1]->get_shape().lens()));
 
-        std::vector<size_t> expected_scales_lens{n * n_blocks_per_col};
-        if(args[2]->get_shape().lens() != expected_scales_lens)
+        const size_t expected_scales_lens = n * n_blocks_per_col;
+        if(args[2]->get_shape().elements() != expected_scales_lens)
             MIGRAPHX_THROW("MatMulNBits: Input scales does not match expected dims: " +
-                           to_string_range(expected_scales_lens) +
+                           to_string(expected_scales_lens) +
                            ". Actual dims: " + to_string_range(args[2]->get_shape().lens()));
 
         if(args.size() > 3)