Merge branch 'bowa_fallback' of https://github.com/NVIDIA/TRTorch int…

…o bowa_fallback
pytorch · Mar 30, 2021 · cfc68ce · cfc68ce
2 parents 459a9b9 + 77b4dc7
commit cfc68ce
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/core/partitioning/partitioning.cpp b/core/partitioning/partitioning.cpp
@@ -106,7 +106,9 @@ void registerSegmentInOutIValues(
   // set inputs ivalues, now supports Tensor/Int to pass argumentes between different segments
   for (auto& input : seg_block.raw_inputs()) {
     TRTORCH_CHECK(ivalues_maps.count(input), "Could not find mini graph input IValue " << input->debugName());
-    if (input->type()->isSubtypeOf(torch::jit::TensorType::get())) {
+    if (input->node()->kind() == torch::jit::prim::Param) {
+      jit_inputs_ivalues.push_back(ivalues_maps[input]);
+    } else if (input->type()->isSubtypeOf(torch::jit::TensorType::get())) {
       jit_inputs_ivalues.push_back(ivalues_maps[input].toTensor());
     } else if (input->type()->isSubtypeOf(torch::jit::IntType::get())) {
       jit_inputs_ivalues.push_back(ivalues_maps[input].toInt());