From 06ef8e85c77b867a11b177fbd2a282672d832980 Mon Sep 17 00:00:00 2001 From: Henry Date: Sun, 24 Nov 2024 07:32:04 +0000 Subject: [PATCH] [Fix] Fix IterBased Loop Training with Faster Resume --- mmengine/runner/loops.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mmengine/runner/loops.py b/mmengine/runner/loops.py index 5a678db7b9..9889681344 100644 --- a/mmengine/runner/loops.py +++ b/mmengine/runner/loops.py @@ -3,6 +3,7 @@ import logging import time from typing import Dict, List, Optional, Sequence, Tuple, Union +from itertools import islice import torch from torch.utils.data import DataLoader @@ -280,8 +281,8 @@ def run(self) -> None: 'that has already been trained', logger='current', level=logging.WARNING) - for _ in range(self._iter): - next(self.dataloader_iterator) + islice(self.dataloader_iterator, 0, self._iter) + while self._iter < self._max_iters and not self.stop_training: self.runner.model.train()