Add PyTorch AMP check (#7917)

* Add PyTorch AMP check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup

* Cleanup

* Cleanup

* Robust for DDP

* Fixes

* Add amp enabled boolean to check_train_batch_size

* Simplify

* space to prefix

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

utils/autobatch.py

@@ -7,15 +7,14 @@ from copy import deepcopy
 
 import numpy as np
 import torch
-from torch.cuda import amp
 
 from utils.general import LOGGER, colorstr
 from utils.torch_utils import profile
 
 
-def check_train_batch_size(model, imgsz=640):
+def check_train_batch_size(model, imgsz=640, amp=True):
     # Check YOLOv5 training batch size
-    with amp.autocast():
+    with torch.cuda.amp.autocast(amp):
         return autobatch(deepcopy(model).train(), imgsz)  # compute optimal batch size