test case for more output types

Author: bowang007

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -259,13 +259,18 @@ def aten_ops_neg(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    input_val = args[0]
+    if (isinstance(input_val, TRTTensor)) and (
+        input_val.dtype == trt.int8 or input_val.dtype == trt.int32
+    ):
+        input_val = cast_trt_tensor(network, input_val, trt.float32, name)
 
     return impl.unary.neg(
         network,
         target,
         SourceIR.ATEN,
         name,
-        args[0],
+        input_val,
     )
 
 

py/torch_tensorrt/dynamo/test_utils.py

@@ -261,6 +261,7 @@ def run_test(
         atol=1e-03,
         precision=torch.float,
         check_dtype=True,
+        output_dtypes=None,
     ):
         mod.eval()
         mod = self.generate_graph(mod, inputs, expected_ops, unexpected_ops, None)
@@ -272,6 +273,7 @@ def run_test(
         interp = TRTInterpreter(
             mod,
             Input.from_tensors(inputs),
+            output_dtypes=output_dtypes,
         )
         super().run_test(
             mod,

tests/py/dynamo/converters/test_neg_aten.py

@@ -11,40 +11,42 @@ class TestNegConverter(DispatchTestCase):
         [
             ("2d_dim_dtype_float", (2, 2), torch.float),
             ("3d_dim_dtype_float", (2, 2, 2), torch.float),
-        
+            ("2d_dim_dtype_half", (2, 2), torch.half),
+            ("3d_dim_dtype_half", (2, 2, 2), torch.half),
         ]
     )
     def test_neg_float(self, _, x, type):
         class neg(nn.Module):
             def forward(self, input):
                 return torch.neg(input)
-            
+
         inputs = [torch.randn(x, dtype=type)]
         self.run_test(
             neg(),
             inputs,
+            precision=type,
             expected_ops={torch.ops.aten.neg.default},
         )
 
     @parameterized.expand(
         [
-            ("2d_dim_dtype_int", (2, 2), torch.int32, 0, 5),
-            ("3d_dim_dtype_int", (2, 2, 2), torch.int32, 0, 5),
+            ("2d_dim_dtype_int32", (2, 2), torch.int32, 0, 5),
+            ("3d_dim_dtype_int32", (2, 2, 2), torch.int32, 0, 5),
         ]
     )
-
     def test_neg_int(self, _, x, type, min, max):
         class neg(nn.Module):
             def forward(self, input):
                 return torch.neg(input)
-            
-        inputs = [torch.randint(min, max, (x), dtype=type)]
 
+        inputs = [torch.randint(min, max, x, dtype=type)]
         self.run_test(
             neg(),
             inputs,
+            output_dtypes=[torch.int32],
             expected_ops={torch.ops.aten.neg.default},
         )
 
+
 if __name__ == "__main__":
     run_tests()