changed enabled parameter

train.py

@@ -37,6 +37,7 @@
 import yaml
 from torch.optim import lr_scheduler
 from tqdm import tqdm
+from torch.cuda import amp  
 
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLOv5 root directory
@@ -352,7 +353,7 @@ def lf(x):
     maps = np.zeros(nc)  # mAP per class
     results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
     scheduler.last_epoch = start_epoch - 1  # do not move
-    scaler = torch.amp.GradScaler(enabled=amp)
+    scaler = amp.GradScaler(enabled=cuda)
     stopper, stop = EarlyStopping(patience=opt.patience), False
     compute_loss = ComputeLoss(model)  # init loss class
     callbacks.run("on_train_start")