Removing the int64 casting to int32 in harness.py since native int64 …

…in supported in TRT10
pytorch · May 7, 2024 · 19049bc · 19049bc
1 parent 6afbb8d
commit 19049bc
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py
@@ -247,10 +247,8 @@ def run_test(
         trt_inputs = inputs
         for num_input in range(num_inputs):
             input = inputs[num_input]
-            if input.dtype in (torch.int64, torch.float64):
-                dtype_32bit = (
-                    torch.int32 if (input.dtype == torch.int64) else torch.float32
-                )
+            if input.dtype is torch.float64:
+                dtype_32bit = torch.float32
                 # should we modify graph here to insert clone nodes?
                 # ideally not required
                 trt_inputs = (