[Enhance] Set torch seed in `worker_init_fn` (#733)

2022-03-16 23:12:57 +08:00 · 2022-03-16 23:12:57 +08:00 · 7856141132
parent aa522f4309
commit 7856141132
1 changed files with 1 additions and 0 deletions
--- 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):