Merge pull request #308 from NVIDIA/fix_num_to_tensor

Fix issue with num_to_tensor evaluator

Author: narendasan

core/conversion/evaluators/prim.cpp

@@ -30,7 +30,7 @@ auto prim_registrations =
                     }})
         .evaluator({torch::jit::prim::NumToTensor,
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
-                      return at::scalar_to_tensor(args.at(n->output(0)).IValue()->toScalar());
+                      return at::scalar_to_tensor(args.at(n->input(0)).IValue()->toScalar());
                     }})
         .evaluator({torch::jit::prim::ListUnpack,
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {

tests/core/conversion/evaluators/test_prim_evaluators.cpp

@@ -37,3 +37,19 @@ TEST(Evaluators, PrimListUnpackEvaluatesCorrectly) {
   ASSERT_TRUE(jit_results[0] == trt_results[0]);
   ASSERT_TRUE(jit_results[1] == trt_results[1]);
 }
+
+TEST(Evaluators, NumToTensorEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : int = prim::Constant[value=3]()
+        %lu.1 : Tensor = prim::NumToTensor(%1)
+        return (%lu.1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}