1 files changed, 10 insertions, 1 deletions
diff --git a/src/gpu/cl/kernels/ClGemmMatrixMultiplyReshapedOnlyRhsMMULKernel.cpp b/src/gpu/cl/kernels/ClGemmMatrixMultiplyReshapedOnlyRhsMMULKernel.cpp
index d67c595449..9a2a4890f3 100644
--- a/src/gpu/cl/kernels/ClGemmMatrixMultiplyReshapedOnlyRhsMMULKernel.cpp
+++ b/src/gpu/cl/kernels/ClGemmMatrixMultiplyReshapedOnlyRhsMMULKernel.cpp
@@ -99,7 +99,16 @@ Status validate_arguments(const ITensorInfo       *src0,
     ARM_COMPUTE_UNUSED(k);
 
     ARM_COMPUTE_RETURN_ERROR_ON(src0->dimension(0) != k);
-    ARM_COMPUTE_RETURN_ERROR_ON(src0->dimension(1) != m);
+
+    // Validate the reinterpreted-as-3D-case
+    if (gemm_info.depth_output_gemm3d != 0)
+    {
+        ARM_COMPUTE_RETURN_ERROR_ON(src0->dimension(1) * src0->dimension(2) != m);
+    }
+    else
+    {
+        ARM_COMPUTE_RETURN_ERROR_ON(src0->dimension(1) != m);
+    }
 
     // Validate the gemm-batched case
     if (src1->num_dimensions() > 2)