From 8972621138a7fd59c95c2fe9645bc959fc5c4ede Mon Sep 17 00:00:00 2001
From: Wei-Sheng Chin <wschin@outlook.com>
Date: Wed, 10 Feb 2021 15:44:25 +0800
Subject: [PATCH] Generate shape-independent graph if any input dimension < 2
 (#6581)

* Throw for non-supported case

* Not to go to shape-dependent branch when seeing unsupported shapes
---
 orttraining/orttraining/core/graph/gradient_builder.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc
index 9126884ff7..9c33a032e3 100644
--- a/orttraining/orttraining/core/graph/gradient_builder.cc
+++ b/orttraining/orttraining/core/graph/gradient_builder.cc
@@ -194,7 +194,8 @@ IMPLEMENT_GRADIENT_BUILDER(GetMatMulGradient) {
     }
   };
 
-  if (A_has_shape && B_has_shape && Y_has_shape) {
+  if (A_has_shape && B_has_shape && Y_has_shape &&
+      A_shape.size() >= 2 && B_shape.size() >= 2) {
     std::vector<AttributeProto> shared_attributes;
     shared_attributes.push_back(MakeAttribute("beta", float(0)));
     AttributeProto transpose_first_input = MakeAttribute("transA", int64_t(1));