Merge pull request #38 from yurujaja/37-dataloader-killed-system-memo…

…ry-utilisation-is-monotonically-increasing-during-training Remove persistent workers. Fix #37
VMarsocci · Sep 16, 2024 · 879e854 · 879e854
2 parents e6ef2e6 + 50e2612
commit 879e854
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/run.py b/run.py
@@ -255,7 +255,8 @@ def main():
             batch_size=cfg.batch_size,  # cfg.dataset["batch"],
             num_workers=cfg.num_workers,
             pin_memory=True,
-            persistent_workers=True,
+            # persistent_workers=True causes memory leak
+            persistent_workers=False,
             worker_init_fn=seed_worker,
             generator=get_generator(cfg.seed),
             drop_last=True,
@@ -268,7 +269,7 @@ def main():
             num_workers=cfg.num_workers,
             pin_memory=True,
             persistent_workers=False,
-            # worker_init_fn=seed_worker,
+            worker_init_fn=seed_worker,
             # generator=g,
             drop_last=False,
             collate_fn=collate_fn,