From e2e95b2d8e8b6be216f4a7c11955d622aff7d043 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 18 Jan 2022 13:52:25 -1000 Subject: [PATCH] TensorRT `assert im.device.type != 'cpu'` on export (#6340) * TensorRT `assert im.device.type != 'cpu'` on export * Update export.py --- export.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/export.py b/export.py index 11fead4a9e1d..2e90b0a1b24c 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 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) if verbose: