From d5651dfa992442cbb698d837266c8ea6c3697b57 Mon Sep 17 00:00:00 2001
From: Jane-Ding <69578198+Jane-Ding@users.noreply.github.com>
Date: Thu, 29 Apr 2021 12:37:05 +0800
Subject: [PATCH] fix win train loader

---
 tools/program.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tools/program.py b/tools/program.py
index db8e44df..7e54a2f8 100755
--- 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:
-- 
GitLab