democat commited on
Commit
f4b0568
·
unverified ·
1 Parent(s): f8722b4

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

Browse files
Files changed (1) hide show
  1. export.py +2 -0
export.py CHANGED
@@ -492,6 +492,8 @@ def run(
492
  # Checks
493
  imgsz *= 2 if len(imgsz) == 1 else 1 # expand
494
  assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
 
 
495
 
496
  # Input
497
  gs = int(max(model.stride)) # grid size (max stride)
 
492
  # Checks
493
  imgsz *= 2 if len(imgsz) == 1 else 1 # expand
494
  assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
495
+ if optimize:
496
+ assert device.type != 'cuda', '--optimize not compatible with cuda devices, i.e. use --device cpu'
497
 
498
  # Input
499
  gs = int(max(model.stride)) # grid size (max stride)