From df838a42dd3c9bd02e1975ae94a1fb855cb94fcf Mon Sep 17 00:00:00 2001
From: Bala-Vignesh-Reddy <reddybalavignesh9979@gmail.com>
Date: Sun, 5 Jan 2025 21:29:50 +0530
Subject: [PATCH] changed enabled parameter

---
 train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 3ebe570b67c4..6ad263722bbf 100644
--- a/train.py
+++ b/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")