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

ctjanuhowski · Sep 8, 2022 · 8080bbb · 8080bbb
1 parent fd43289
commit 8080bbb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/export.py b/export.py
@@ -611,6 +611,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)