|
15 | 15 | # Importing supported Backends
|
16 | 16 | import torch
|
17 | 17 | import torch_tensorrt as torchtrt
|
18 |
| -from torch_tensorrt.fx.lower import lower_to_trt |
| 18 | +from torch_tensorrt.fx.lower import compile |
19 | 19 | from torch_tensorrt.fx.utils import LowerPrecision
|
20 | 20 |
|
21 | 21 | import tensorrt as trt
|
@@ -89,7 +89,7 @@ def run_torch_tensorrt(model, input_tensors, params, precision, truncate_long_an
|
89 | 89 | if precision == "int8":
|
90 | 90 | compile_settings.update({"calib": params.get("calibration_cache")})
|
91 | 91 |
|
92 |
| - with torchtrt.logging.errors(): |
| 92 | + with torchtrt.logging.debug(): |
93 | 93 | model = torchtrt.compile(model, **compile_settings)
|
94 | 94 |
|
95 | 95 | iters = params.get("iterations", 20)
|
@@ -123,7 +123,7 @@ def run_fx2trt(model, input_tensors, params, precision, batch_size):
|
123 | 123 | model.half()
|
124 | 124 | input_tensors = [tensor.half() for tensor in input_tensors]
|
125 | 125 | # Run lowering eager mode benchmark
|
126 |
| - model = lower_to_trt( |
| 126 | + model = compile( |
127 | 127 | model,
|
128 | 128 | input_tensors,
|
129 | 129 | max_batch_size=batch_size,
|
@@ -430,7 +430,8 @@ def load_model(params):
|
430 | 430 | print("Model Summary: ", model_name)
|
431 | 431 | summary = pd.DataFrame(results)
|
432 | 432 | print(summary)
|
433 |
| - with open(args.report, "w") as file: |
434 |
| - file.write("Model Summary: " + model_name + "\n") |
435 |
| - file.write(summary.to_string()) |
436 |
| - file.close() |
| 433 | + if args.report: |
| 434 | + with open(args.report, "w") as file: |
| 435 | + file.write("Model Summary: " + model_name + "\n") |
| 436 | + file.write(summary.to_string()) |
| 437 | + file.close() |
0 commit comments