Skip to content

Commit 29ce1a7

Browse files
authored
Use torch_tensorrt.Device instead of torch.device in trt compile (#8051)
Fixes #8050 ### Description A few sentences describing the changes proposed in this pull request. ### Types of changes <!--- Put an `x` in all the boxes that apply, and remove the not applicable items --> - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [ ] Breaking change (fix or new feature that would cause existing functionality to change). - [ ] New tests added to cover the changes. - [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [ ] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
1 parent b6d6d77 commit 29ce1a7

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

monai/networks/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -851,7 +851,7 @@ def _onnx_trt_compile(
851851
# wrap the serialized TensorRT engine back to a TorchScript module.
852852
trt_model = torch_tensorrt.ts.embed_engine_in_new_module(
853853
f.getvalue(),
854-
device=torch.device(f"cuda:{device}"),
854+
device=torch_tensorrt.Device(f"cuda:{device}"),
855855
input_binding_names=input_names,
856856
output_binding_names=output_names,
857857
)

0 commit comments

Comments
 (0)