chore: Update for random input

Author: keehyuna

py/torch_tensorrt/_Input.py

@@ -352,9 +352,20 @@ def example_tensor(
                 )
             else:
                 if isinstance(self.shape, tuple):
-                    return torch.rand(self.shape).to(
-                        dtype=self.dtype.to(torch.dtype, use_default=True)
-                    )
+                    if self.dtype in [dtype.u8, dtype.i8, dtype.i32, dtype.i64]:
+                        type = self.dtype.to(torch.dtype, use_default=True)
+                        return torch.randint(
+                            torch.iinfo(type).min,
+                            torch.iinfo(type).max,
+                            self.shape,
+                            dtype=type,
+                        )
+                    elif self.dtype == dtype.b:
+                        return torch.rand(self.shape) < 0.5
+                    else:
+                        return torch.rand(self.shape).to(
+                            dtype=self.dtype.to(torch.dtype, use_default=True)
+                        )
                 else:
                     RuntimeError(
                         f"Input shape is dynamic but shapes are not provided as sequence (found: {self.shape})"
@@ -380,6 +391,8 @@ def example_tensor(
                             self.shape[optimization_profile_field],
                             dtype=type,
                         )
+                    elif self.dtype == dtype.b:
+                        return torch.rand(self.shape[optimization_profile_field]) < 0.5
                     else:
                         return torch.rand(self.shape[optimization_profile_field]).to(
                             dtype=self.dtype.to(torch.dtype, use_default=True)

tests/py/dynamo/conversion/harness.py

@@ -319,7 +319,7 @@ def run_test_compare_tensor_attributes_only(
         expected_ops,
         comparators: List[Tuple[Callable, List]],
         precision=torch.float,
-        output_dtypes=None,
+        check_dtype=True,
         use_dynamo_tracer=False,
         enable_passes=False,
     ):
@@ -337,6 +337,15 @@ def run_test_compare_tensor_attributes_only(
             debug=True,
         )
 
+        output_dtypes = None
+        if check_dtype:
+            output_dtypes = infer_module_output_dtypes(
+                mod,
+                [Input.from_tensor(i) for i in inputs],
+                compilation_settings.device,
+                truncate_double=compilation_settings.truncate_double,
+            )
+
         interp = TRTInterpreter(
             mod,
             Input.from_tensors(inputs),
@@ -393,4 +402,12 @@ def run_test_with_dynamic_shape(
         inputs_max = [spec.example_tensor("max_shape") for spec in input_specs]
         if not use_example_tensors:
             inputs_max = [spec.torch_tensor for spec in input_specs]
-        super().run_test(mod, inputs_max, interp, rtol, atol, check_dtype=check_dtype, pyt_inputs=pyt_inputs)
+        super().run_test(
+            mod,
+            inputs_max,
+            interp,
+            rtol,
+            atol,
+            check_dtype=check_dtype,
+            pyt_inputs=pyt_inputs,
+        )