diff --git a/mmdet/datasets/loader/sampler.py b/mmdet/datasets/loader/sampler.py index 1e454b741c8e2dbee10ba25b928a13ba1da18392..f45ba1325d5ec235d28ec92f460820a3adfee7cf 100644 --- a/mmdet/datasets/loader/sampler.py +++ b/mmdet/datasets/loader/sampler.py @@ -4,7 +4,7 @@ import math import torch import numpy as np -from torch.distributed import get_world_size, get_rank +from mmcv.runner.utils import get_dist_info from torch.utils.data import Sampler from torch.utils.data import DistributedSampler as _DistributedSampler @@ -95,10 +95,11 @@ class DistributedGroupSampler(Sampler): samples_per_gpu=1, num_replicas=None, rank=None): + _rank, _num_replicas = get_dist_info() if num_replicas is None: - num_replicas = get_world_size() + num_replicas = _num_replicas if rank is None: - rank = get_rank() + rank = _rank self.dataset = dataset self.samples_per_gpu = samples_per_gpu self.num_replicas = num_replicas