diff --git a/utils/general.py b/utils/general.py index 91b13f8..842f28c 100755 --- a/utils/general.py +++ b/utils/general.py @@ -535,8 +535,8 @@ def check_amp(model): prefix = colorstr('AMP: ') device = next(model.parameters()).device # get model device - if device.type == 'cpu': - return False # AMP disabled on CPU + if device.type in ('cpu', 'mps'): + return False # AMP only used on CUDA devices f = ROOT / 'data' / 'images' / 'bus.jpg' # image to check im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if check_online() else np.ones((640, 640, 3)) try: