fix memory bug

bubbliiiing · Jul 2, 2022 · 0a289c0 · 0a289c0
1 parent e618179
commit 0a289c0
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/utils/utils_fit.py b/utils/utils_fit.py
@@ -21,8 +21,8 @@ def fit_one_epoch(model_train, model, yolo_loss, loss_history, eval_callback, op
         images, targets = batch[0], batch[1]
         with torch.no_grad():
             if cuda:
-                images  = images.cuda()
-                targets = [ann.cuda() for ann in targets]
+                images  = images.cuda(local_rank)
+                targets = [ann.cuda(local_rank) for ann in targets]
         #----------------------#
         #   清零梯度
         #----------------------#
@@ -91,8 +91,8 @@ def fit_one_epoch(model_train, model, yolo_loss, loss_history, eval_callback, op
         images, targets = batch[0], batch[1]
         with torch.no_grad():
             if cuda:
-                images  = images.cuda()
-                targets = [ann.cuda() for ann in targets]
+                images  = images.cuda(local_rank)
+                targets = [ann.cuda(local_rank) for ann in targets]
             #----------------------#
             #   清零梯度
             #----------------------#