feat(prim::NumToTensor): Implement evaluator for NumToTensor

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Jun 1, 2020 · 60df888 · 60df888
1 parent 17099fa
commit 60df888
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/core/conversion/evaluators/prim.cpp b/core/conversion/evaluators/prim.cpp
@@ -5,6 +5,7 @@
 #include "ATen/core/List.h"
 #include "ATen/core/stack.h"
 #include "c10/util/intrusive_ptr.h"
+#include "torch/torch.h"
 
 #include "core/conversion/evaluators/evaluators.h"
 
@@ -23,6 +24,11 @@ auto prim_registrations = RegisterNodeEvaluators()
             }
             return torch::jit::toIValue(n->output());
         }
+    }).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());
+        }
     }).evaluator({
         torch::jit::prim::ListConstruct,
         [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {