Merge pull request #1095 from mfeliz-cruise/michael.feliz/int_scalar_mul

fix: support int tensor * int scaler in aten::mul

Author: peri044

core/conversion/converters/impl/element_wise.cpp

@@ -425,8 +425,7 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     // TODO: Remove with functionalization
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto otherScalar = args[1].unwrapToScalar().to<float>();
-                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
                     auto mul =
                         add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
                     TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n);

tests/core/conversion/converters/test_element_wise.cpp

@@ -11,7 +11,8 @@ void pointwise_test_helper(
     bool dynamicInput = false,
     std::vector<int64_t> shape1 = {5},
     std::vector<int64_t> shape2 = {5},
-    bool negative_input = false) {
+    bool negative_input = false,
+    bool int_tensors = false) {
   auto g = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(graph_ir, g.get());
 
@@ -26,6 +27,11 @@ void pointwise_test_helper(
   if (!singleInput) {
     torch_inputs.push_back(at::randint(1, 5, shape2, {at::kCUDA}));
   }
+  if(int_tensors){
+    for(size_t i = 0UL; i < torch_inputs.size(); ++i){
+      torch_inputs[i] = torch_inputs[i].to(at::kInt);
+    }
+  }
   auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
   auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, torch_inputs);
 
@@ -126,6 +132,15 @@ TEST(Converters, ATenMulWithScalarConvertsCorrectly) {
   pointwise_test_helper(graph, true);
 }
 
+TEST(Converters, ATenMulWithIntScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %scalar : int = prim::Constant[value=2]()
+        %1 : Tensor = aten::mul(%0, %scalar)
+        return (%1))IR";
+  pointwise_test_helper(graph, true, false, {5}, {5}, false, true);
+}
+
 TEST(Converters, ATenDivConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor, %1 : Tensor):