From 571c5dcee325a78cd8f02e2ee8ebb17e1df2ac75 Mon Sep 17 00:00:00 2001
From: Jacky <18255193+kthui@users.noreply.github.com>
Date: Wed, 16 Oct 2024 09:41:55 -0700
Subject: [PATCH] Allow missing ensemble output(s) (#399)

---
 src/ensemble_scheduler/ensemble_scheduler.cc | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/ensemble_scheduler/ensemble_scheduler.cc b/src/ensemble_scheduler/ensemble_scheduler.cc
index b16567dd7..609c1ff28 100644
--- a/src/ensemble_scheduler/ensemble_scheduler.cc
+++ b/src/ensemble_scheduler/ensemble_scheduler.cc
@@ -1200,11 +1200,6 @@ EnsembleContext::CheckAndSetEnsembleOutput(
         ready = false;
         break;
       }
-      // check if the output is provided
-      else if (tensor[iteration_count].data_ == nullptr) {
-        ready = false;
-        break;
-      }
     }
   }
   if (!ready) {
@@ -1224,6 +1219,12 @@ EnsembleContext::CheckAndSetEnsembleOutput(
     auto& tensor_data = tensor_data_[output_pair.first];
     auto& tensor = tensor_data.tensor_[iteration_count];
 
+    if (tensor.data_ == nullptr) {
+      LOG_VERBOSE(1) << "Composing models did not output tensor "
+                     << output_pair.first;
+      continue;
+    }
+
     auto shape = ReshapeTensorDims(
         output_pair.second, (lrequest->BatchSize() != 0),
         tensor_data.batch_size_, tensor.data_->OriginalShape());