From cfb5fb7e9f428ca7442bb9b4a8d6c238f6f94360 Mon Sep 17 00:00:00 2001
From: Ubuntu
 <OrtTrainingDev3@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
Date: Fri, 22 Jan 2021 18:04:20 +0000
Subject: [PATCH 1/3] Continue memory planning when unknown shape tensor is
 encountered.

---
 onnxruntime/core/framework/session_state.cc | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc
index 191526e455549..12b4f73535103 100644
--- a/onnxruntime/core/framework/session_state.cc
+++ b/onnxruntime/core/framework/session_state.cc
@@ -389,6 +389,7 @@ void TryCalculateSizeFromResolvedShape(int ml_value_idx, std::unordered_map<int,
 
 }  // namespace
 
+// If this function fails NO TRAINING will take place, hence lets ONLY FAIL and stop training where warranted, example SIZE overflow.
 Status SessionState::GeneratePatternGroupCache(const std::vector<std::reference_wrapper<const TensorShape>>& input_shape,
                                                const std::vector<int>& feed_mlvalue_idxs,
                                                MemoryPatternGroup* output,
@@ -425,12 +426,14 @@ Status SessionState::GeneratePatternGroupCache(const std::vector<std::reference_
       auto* arg = node->OutputDefs()[i];
       size_t is_resolved = 0;
       std::vector<int64_t> resolved_shape;
-      ORT_RETURN_IF_ERROR(TryResolveShape(arg, map, is_resolved, resolved_shape));
 
-      // Store all valid resolved shapes. They will be queried in, for example,
-      // Recv operator to bypass the dependency of output shapes on inputs.
-      if (is_resolved != 0) {
-        resolved_shapes[ml_value_idx] = resolved_shape;
+      // Tensors whose shape cannot be resolved statically will be allocated at runtime.
+      if (TryResolveShape(arg, map, is_resolved, resolved_shape).IsOK()) {
+        // Store all valid resolved shapes. They will be queried in, for example,
+        // Recv operator to bypass the dependency of output shapes on inputs.
+        if (is_resolved != 0) {
+          resolved_shapes[ml_value_idx] = resolved_shape;
+        }
       }
     }
   }

From 377cb95e0b089649f5ca676f5c38f7a49224cb1e Mon Sep 17 00:00:00 2001
From: Ubuntu
 <OrtTrainingDev3@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
Date: Fri, 22 Jan 2021 18:21:07 +0000
Subject: [PATCH 2/3] comment.

---
 onnxruntime/core/framework/session_state.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc
index 12b4f73535103..a11120720295e 100644
--- a/onnxruntime/core/framework/session_state.cc
+++ b/onnxruntime/core/framework/session_state.cc
@@ -389,7 +389,7 @@ void TryCalculateSizeFromResolvedShape(int ml_value_idx, std::unordered_map<int,
 
 }  // namespace
 
-// If this function fails NO TRAINING will take place, hence lets ONLY FAIL and stop training where warranted, example SIZE overflow.
+// If this function fails NO memory planning will take place, hence lets ONLY FAIL and stop training where warranted, example SIZE overflow.
 Status SessionState::GeneratePatternGroupCache(const std::vector<std::reference_wrapper<const TensorShape>>& input_shape,
                                                const std::vector<int>& feed_mlvalue_idxs,
                                                MemoryPatternGroup* output,

From 5386c95b15c74d351535946b5706423fa36ab4b4 Mon Sep 17 00:00:00 2001
From: Ubuntu
 <OrtTrainingDev3@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
Date: Fri, 22 Jan 2021 19:19:32 +0000
Subject: [PATCH 3/3] PR feedback.

---
 onnxruntime/core/framework/session_state.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc
index a11120720295e..13f5ddf4f7d8c 100644
--- a/onnxruntime/core/framework/session_state.cc
+++ b/onnxruntime/core/framework/session_state.cc
@@ -434,6 +434,9 @@ Status SessionState::GeneratePatternGroupCache(const std::vector<std::reference_
         if (is_resolved != 0) {
           resolved_shapes[ml_value_idx] = resolved_shape;
         }
+      } else {
+        LOGS(logger_, INFO) << "[Static memory planning] Could not resolve shape for tensor with ML index "
+                            << ml_value_idx << ", will allocate dynamically.";
       }
     }
   }