From f4b05680f89795658e1c898a28ff51edbf22a63b Mon Sep 17 00:00:00 2001 From: Colin Wong Date: Fri, 15 Jul 2022 09:01:01 -0500 Subject: [PATCH] Assert `--optimize` not used with cuda device (#8569) --- export.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/export.py b/export.py index 623844ff3531..9868fcae95c3 100644 --- a/export.py +++ b/export.py @@ -492,6 +492,8 @@ def run( # Checks imgsz *= 2 if len(imgsz) == 1 else 1 # expand assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}' + if optimize: + assert device.type != 'cuda', '--optimize not compatible with cuda devices, i.e. use --device cpu' # Input gs = int(max(model.stride)) # grid size (max stride)