From 630f9c45f70a1810ff8e01c3c6a3da65ed30a7a3 Mon Sep 17 00:00:00 2001 From: Bo Wang Date: Thu, 28 Apr 2022 15:11:06 -0700 Subject: [PATCH] fix: support dict type for input in shape analysis Signed-off-by: Bo Wang --- core/partitioning/shape_analysis.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/core/partitioning/shape_analysis.cpp b/core/partitioning/shape_analysis.cpp index e773b8f321..a899b3313e 100644 --- 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()