From 0dfafaf39b263cdb8d7e637ef37df825ebe72576 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 17 Jun 2023 01:07:24 +0200
Subject: [PATCH] Fix `batch=-1` for exports (#3217)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 ultralytics/yolo/engine/model.py    | 2 +-
 ultralytics/yolo/engine/trainer.py  | 2 +-
 ultralytics/yolo/utils/autobatch.py | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py
index 6d3168b..014c18e 100644
--- a/ultralytics/yolo/engine/model.py
+++ b/ultralytics/yolo/engine/model.py
@@ -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
diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py
index c3f8fcd..679c545 100644
--- a/ultralytics/yolo/engine/trainer.py
+++ b/ultralytics/yolo/engine/trainer.py
@@ -229,7 +229,7 @@ class BaseTrainer:
         # Batch size
         if self.batch_size == -1:
             if RANK == -1:  # single-GPU only, estimate best batch size
-                self.batch_size = check_train_batch_size(self.model, self.args.imgsz, self.amp)
+                self.args.batch = self.batch_size = check_train_batch_size(self.model, self.args.imgsz, self.amp)
             else:
                 SyntaxError('batch=-1 to use AutoBatch is only available in Single-GPU training. '
                             'Please pass a valid batch size value for Multi-GPU DDP training, i.e. batch=16')
diff --git a/ultralytics/yolo/utils/autobatch.py b/ultralytics/yolo/utils/autobatch.py
index e845579..0645f81 100644
--- a/ultralytics/yolo/utils/autobatch.py
+++ b/ultralytics/yolo/utils/autobatch.py
@@ -8,7 +8,7 @@ from copy import deepcopy
 import numpy as np
 import torch
 
-from ultralytics.yolo.utils import LOGGER, colorstr
+from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, colorstr
 from ultralytics.yolo.utils.torch_utils import profile
 
 
@@ -29,7 +29,7 @@ def check_train_batch_size(model, imgsz=640, amp=True):
         return autobatch(deepcopy(model).train(), imgsz)  # compute optimal batch size
 
 
-def autobatch(model, imgsz=640, fraction=0.67, batch_size=16):
+def autobatch(model, imgsz=640, fraction=0.67, batch_size=DEFAULT_CFG.batch):
     """
     Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.