Assert --optimize not used with cuda device (#8569)

ultralytics · Jul 15, 2022 · f4b0568 · f4b0568
1 parent f8722b4
commit f4b0568
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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)