File tree Expand file tree Collapse file tree 1 file changed +4
-0
lines changed
tests/py/dynamo/converters Expand file tree Collapse file tree 1 file changed +4
-0
lines changed Original file line number Diff line number Diff line change @@ -266,6 +266,7 @@ def run_test(
266
266
precision = torch .float ,
267
267
check_dtype = True ,
268
268
disable_passes = False ,
269
+ output_dtypes = None ,
269
270
):
270
271
mod .eval ()
271
272
mod = self .generate_graph (
@@ -284,6 +285,7 @@ def run_test(
284
285
interp = TRTInterpreter (
285
286
mod ,
286
287
Input .from_tensors (inputs ),
288
+ output_dtypes = output_dtypes ,
287
289
)
288
290
super ().run_test (
289
291
mod ,
@@ -306,6 +308,7 @@ def run_test_with_dynamic_shape(
306
308
rtol = 1e-03 ,
307
309
atol = 1e-03 ,
308
310
disable_passes = False ,
311
+ output_dtypes = None ,
309
312
):
310
313
mod .eval ()
311
314
inputs = [spec .example_tensor ("opt_shape" ) for spec in input_specs ]
@@ -321,6 +324,7 @@ def run_test_with_dynamic_shape(
321
324
interp = TRTInterpreter (
322
325
mod ,
323
326
input_specs ,
327
+ output_dtypes = output_dtypes ,
324
328
)
325
329
# Since the lowering is based on optimal shape. We need to test with
326
330
# different shape(for ex. max shape) for testing dynamic shape
You can’t perform that action at this time.
0 commit comments