From fca53cef9be1021022c98a8d3286bde4166117f4 Mon Sep 17 00:00:00 2001
From: inocsin <vcheungyi@163.com>
Date: Wed, 20 Oct 2021 08:54:46 +0000
Subject: [PATCH] feat: handle scalar type of size [] in shape_analysis

Signed-off-by: inocsin <vcheungyi@163.com>
---
 core/partitioning/shape_analysis.cpp | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/core/partitioning/shape_analysis.cpp b/core/partitioning/shape_analysis.cpp
index 713fd76e46..3657cecdaf 100644
--- a/core/partitioning/shape_analysis.cpp
+++ b/core/partitioning/shape_analysis.cpp
@@ -124,7 +124,12 @@ void getSegmentsOutputByRunning(
       if (dtype == c10::nullopt) {
         TRTORCH_THROW_ERROR("Unsupported input data type " << ivalues_maps[i].toTensor().dtype());
       }
-      input_shapes.push_back(util::toVec(util::toDims(ivalues_maps[i].toTensor().sizes())));
+      if (ivalues_maps[i].toTensor().sizes().size() == 0) {
+        // handle Scalar types, which has sizes of []
+        input_shapes.push_back(util::toVec(util::toDims(c10::List<long int>({1}))));
+      } else {
+        input_shapes.push_back(util::toVec(util::toDims(ivalues_maps[i].toTensor().sizes())));
+      }
       input_types.push_back(ivalues_maps[i].toTensor().scalar_type());
     }
   }