diff --git a/tests/py/dynamo/converters/harness.py b/tests/py/dynamo/converters/harness.py index f6ff25fb77..c8ea5bb5c0 100644 --- a/tests/py/dynamo/converters/harness.py +++ b/tests/py/dynamo/converters/harness.py @@ -266,6 +266,7 @@ def run_test( precision=torch.float, check_dtype=True, disable_passes=False, + output_dtypes=None, ): mod.eval() mod = self.generate_graph( @@ -284,6 +285,7 @@ def run_test( interp = TRTInterpreter( mod, Input.from_tensors(inputs), + output_dtypes=output_dtypes, ) super().run_test( mod, @@ -306,6 +308,7 @@ def run_test_with_dynamic_shape( rtol=1e-03, atol=1e-03, disable_passes=False, + output_dtypes=None, ): mod.eval() inputs = [spec.example_tensor("opt_shape") for spec in input_specs] @@ -321,6 +324,7 @@ def run_test_with_dynamic_shape( interp = TRTInterpreter( mod, input_specs, + output_dtypes=output_dtypes, ) # Since the lowering is based on optimal shape. We need to test with # different shape(for ex. max shape) for testing dynamic shape