fix: support dict type for input in shape analysis

Signed-off-by: Bo Wang <[email protected]>
pytorch · Apr 28, 2022 · 630f9c4 · 630f9c4
1 parent 1c294fa
commit 630f9c4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/core/partitioning/shape_analysis.cpp b/core/partitioning/shape_analysis.cpp
@@ -84,6 +84,8 @@ void getSegmentsOutputByRunning(
       jit_inputs_ivalues.push_back(ivalues_maps[input].toTuple());
     } else if (input->type()->kind() == torch::jit::TypeKind::NumberType) {
       jit_inputs_ivalues.push_back(ivalues_maps[input].toScalar());
+    } else if (input->type()->kind() == torch::jit::TypeKind::DictType) {
+      jit_inputs_ivalues.push_back(ivalues_maps[input].toGenericDict());
     } else {
       TORCHTRT_THROW_ERROR(
           "Expected to find type " << input->type()->str() << " for value " << input->debugName()