Merge pull request #964 from Njuapp/roberta_fix

Fix roberta conversion bugs

Author: narendasan

core/conversion/converters/impl/cast.cpp

@@ -18,7 +18,14 @@ auto cast_registrations TORCHTRT_UNUSED =
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
                auto output_dtype = args[1].unwrapToScalar().to<int64_t>();
-               auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
+               auto scalar_dtype = static_cast<at::ScalarType>(output_dtype);
+               nvinfer1::DataType trt_dtype;
+               if (scalar_dtype == at::kLong) {
+                 LOG_WARNING("Truncating aten::to output type from at::kLong to at::kInt");
+                 trt_dtype = nvinfer1::DataType::kINT32;
+               } else {
+                 trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
+               }
                auto casted_itensor = castITensor(ctx, self, trt_dtype);
                auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
                LOG_DEBUG("[aten::to.dtype] Output tensor shape: " << output->getDimensions());
@@ -33,9 +40,14 @@ auto cast_registrations TORCHTRT_UNUSED =
                // later shape analysis phase of fallback
                auto self = args[0].ITensorOrFreeze(ctx);
                auto output_dtype = args[2].unwrapToScalar().to<int64_t>();
-
-               auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
-
+               auto scalar_dtype = static_cast<at::ScalarType>(output_dtype);
+               nvinfer1::DataType trt_dtype;
+               if (scalar_dtype == at::kLong) {
+                 LOG_WARNING("Truncating aten::to output type from at::kLong to at::kInt");
+                 trt_dtype = nvinfer1::DataType::kINT32;
+               } else {
+                 trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
+               }
                auto casted_itensor = castITensor(ctx, self, trt_dtype);
                auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
                LOG_DEBUG("[aten::to.device] Output tensor shape: " << output->getDimensions());

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,13 @@ 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 (self->getType() == nvinfer1::DataType::kFLOAT || self->getType() == nvinfer1::DataType::kHALF) {
+                      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,