Merge pull request #972 from NVIDIA/maskrcnn

fix: fix the bug that introduces kLong Tensor in prim::NumToTensor

Author: narendasan

core/conversion/evaluators/eval_util.cpp

@@ -119,6 +119,39 @@ void checkSequenceSize(int64_t n, int64_t dim, int64_t seq_size) {
   }
 }
 
+// TODO: Conditionally enable truncation based on user setting
+at::Tensor scalar_to_tensor(const at::Scalar& s, const at::Device device = at::kCPU) {
+  // This function is basically same with the one in
+  // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/ScalarOps.h, what different here is that Int and Float
+  // won't be upgraded to kDouble or kLong since we don't support these 2 types in conversion
+  if (device == at::kCPU) {
+    if (s.isFloatingPoint()) {
+      LOG_WARNING("Unable to process input type of at::kDouble, truncate type to at::kFloat in scalar_to_tensor_util ");
+      return at::detail::scalar_tensor_static(s, at::kFloat, at::kCPU);
+    } else if (s.isComplex()) {
+      return at::detail::scalar_tensor_static(s, at::kComplexDouble, at::kCPU);
+    } else if (s.isBoolean()) {
+      return at::detail::scalar_tensor_static(s, at::kBool, at::kCPU);
+    } else {
+      AT_ASSERT(s.isIntegral(false));
+      LOG_WARNING("Unable to process input type of at::kLong, truncate type to at::kInt in scalar_to_tensor_util ");
+      return at::detail::scalar_tensor_static(s, at::kInt, at::kCPU);
+    }
+  }
+  if (s.isFloatingPoint()) {
+    LOG_WARNING("Unable to process input type of at::kDouble, truncate type to at::kFloat in scalar_to_tensor_util ");
+    return at::scalar_tensor(s, at::device(device).dtype(at::kFloat));
+  } else if (s.isBoolean()) {
+    return at::scalar_tensor(s, at::device(device).dtype(at::kBool));
+  } else if (s.isComplex()) {
+    return at::scalar_tensor(s, at::device(device).dtype(at::kComplexDouble));
+  } else {
+    AT_ASSERT(s.isIntegral(false));
+    LOG_WARNING("Unable to process input type of at::kLong, truncate type to at::kInt in scalar_to_tensor_util ");
+    return at::scalar_tensor(s, at::device(device).dtype(at::kInt));
+  }
+}
+
 template <typename DTYPE>
 void storeLastDimension(
     char* data,

core/conversion/evaluators/eval_util.h

@@ -13,6 +13,8 @@ at::Tensor createTensorFromList(
     const torch::jit::IValue& dtype,
     const torch::jit::IValue& device);
 
+at::Tensor scalar_to_tensor(const at::Scalar& s, const at::Device device = at::kCPU);
+
 } // namespace evaluators
 } // namespace conversion
 } // namespace core

core/conversion/evaluators/prim.cpp

@@ -31,7 +31,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->input(0)).IValue()->toScalar());
+                      return evaluators::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> {