fix win train loader

d5651dfa · Jane-Ding · littletomatodonkey · d63431a6 · d5651dfa
Commit d5651dfa authored Apr 29, 2021 by Jane-Ding Committed by littletomatodonkey Apr 29, 2021
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

tools/program.py tools/program.py +5 -1

No files found.
--- a/tools/program.py
+++ b/tools/program.py
@@ -199,8 +199,12 @@ def train(config,
        train_reader_cost = 0.0
        batch_sum = 0
        batch_start = time.time()
-        for idx, batch in enumerate(train_dataloader()):
+        max_iter = len(train_dataloader) - 1 if platform.system(
+        ) == "Windows" else len(train_dataloader)
+        for idx, batch in enumerate(train_dataloader):
            train_reader_cost += time.time() - batch_start
+            if idx >= max_iter:
+                break
            lr = optimizer.get_lr()
            images = batch[0]
            if use_srn: