Fix batch=-1 for exports (#3217)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-06-17 01:07:24 +02:00 committed by GitHub
parent d69a1e8046
commit 0dfafaf39b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 4 additions and 4 deletions

View file

@ -335,7 +335,7 @@ class YOLO:
overrides['mode'] = 'export'
if overrides.get('imgsz') is None:
overrides['imgsz'] = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
if overrides.get('batch') is None:
if 'batch' not in kwargs:
overrides['batch'] = 1 # default to 1 if not modified
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.task = self.task