Fix roberta conversion bugs

Signed-off-by: Cheng Hang <[email protected]>

Author: Njuapp

core/conversion/converters/impl/cumsum.cpp

@@ -48,7 +48,8 @@ auto cumsum_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pat
        auto data = iterator->getOutput(0);
        auto newDims = data->getDimensions();
 
-       torch::Tensor zeroValue = at::full(util::toVec(newDims), 0, torch::kFloat32);
+       torch::Tensor zeroValue =
+           at::full(util::toVec(newDims), 0, torch_tensorrt::core::util::TRTDataTypeToScalarType(in->getType()));
        auto zeroTensor = tensor_to_const(ctx, zeroValue);
        auto runningSum = loop->addRecurrence(*zeroTensor);
        auto runningSumTensor = runningSum->getOutput(0);

core/conversion/converters/impl/element_wise.cpp

@@ -484,8 +484,12 @@ auto element_wise_registrations TORCHTRT_UNUSED =
         .pattern({"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto scalar = args[1].unwrapToScalar().to<float>();
-                    auto scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                    auto scalar = args[1].unwrapToScalar();
+                    nvinfer1::ITensor* scalar_tensor;
+                    if (scalar.isFloatingPoint())
+                      scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to<float>()}));
+                    else
+                      scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to<int>()}));
                     auto equal = add_elementwise(
                         ctx,
                         nvinfer1::ElementWiseOperation::kEQUAL,

core/util/trt_util.cpp

@@ -238,6 +238,7 @@ const std::unordered_map<at::ScalarType, nvinfer1::DataType>& get_at_trt_type_ma
       {at::kFloat, nvinfer1::DataType::kFLOAT},
       {at::kHalf, nvinfer1::DataType::kHALF},
       {at::kInt, nvinfer1::DataType::kINT32},
+      {at::kLong, nvinfer1::DataType::kINT32},
       {at::kChar, nvinfer1::DataType::kINT8},
       {at::kBool, nvinfer1::DataType::kBOOL}};
   return at_trt_type_map;