[Enhance] Set torch seed in worker_init_fn (open-mmlab#733)

mzr1996 · Mar 16, 2022 · 7f61b06 · 7f61b06
1 parent 725978a
commit 7f61b06
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mmcls/datasets/builder.py b/mmcls/datasets/builder.py
@@ -156,6 +156,7 @@ def worker_init_fn(worker_id, num_workers, rank, seed):
     worker_seed = num_workers * rank + worker_id + seed
     np.random.seed(worker_seed)
     random.seed(worker_seed)
+    torch.manual_seed(worker_seed)
 
 
 def build_sampler(cfg, default_args=None):