batchmatmul_benchmarks.patch

diff --git a/integration_tests/benchmarks/matmul_benchmarks.ts b/integration_tests/benchmarks/matmul_benchmarks.ts
index 8f67ed9..5e242fb 100644
--- a/integration_tests/benchmarks/matmul_benchmarks.ts
+++ b/integration_tests/benchmarks/matmul_benchmarks.ts
@@ -43,8 +43,8 @@ export class MatmulGPUBenchmark implements BenchmarkTest {
   async run(size: number): Promise<number> {
     tf.setBackend('webgl');
 
-    const a: tf.Tensor2D = tf.randomNormal([size, size]);
-    const b: tf.Tensor2D = tf.randomNormal([size, size]);
+    const a: tf.Tensor3D = tf.randomNormal([3, size, size]);
+    const b: tf.Tensor3D = tf.randomNormal([3, size, size]);
 
     const benchmark = () => tf.matMul(a, b);