From 2a06ad046dd869e0463daf1c5d478bad509f3681 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 18 Jan 2022 10:58:39 -1000 Subject: [PATCH 1/2] TensorRT `assert im.device.type != 'cpu'` on export --- export.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/export.py b/export.py index 11fead4a9e1d..bd28b1f7672c 100644 --- a/export.py +++ b/export.py @@ -184,9 +184,10 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0 export_onnx(model, im, file, 13, train, False, simplify) # opset 13 onnx = file.with_suffix('.onnx') - assert onnx.exists(), f'failed to export ONNX file: {onnx}' LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') + assert im.device.type != 'cpu', 'export must be done on GPU, i.e. try `python export.py --device 0`' + assert onnx.exists(), f'failed to export ONNX file: {onnx}' f = file.with_suffix('.engine') # TensorRT engine file logger = trt.Logger(trt.Logger.INFO) if verbose: From dc679e8bcd3cba506eff90db51da882816f2fa44 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 18 Jan 2022 13:48:21 -1000 Subject: [PATCH 2/2] Update export.py --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index bd28b1f7672c..2e90b0a1b24c 100644 --- a/export.py +++ b/export.py @@ -186,7 +186,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F onnx = file.with_suffix('.onnx') LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') - assert im.device.type != 'cpu', 'export must be done on GPU, i.e. try `python export.py --device 0`' + assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' assert onnx.exists(), f'failed to export ONNX file: {onnx}' f = file.with_suffix('.engine') # TensorRT engine file logger = trt.Logger(trt.Logger.INFO)