Commit b64d4387 authored by Kai Chen's avatar Kai Chen
Browse files

make compatible with pytorch 0.4.1 and 1.0

parent 076cdd6c
......@@ -33,6 +33,9 @@ class MMDistributedDataParallel(nn.Module):
self._dist_broadcast_coalesced(module_states,
self.broadcast_bucket_size)
if self.broadcast_buffers:
if torch.__version__ < '1.0':
buffers = [b.data for b in self.module._all_buffers()]
else:
buffers = [b.data for b in self.module.buffers()]
if len(buffers) > 0:
self._dist_broadcast_coalesced(buffers,
......
......@@ -5,6 +5,7 @@ from getpass import getuser
from socket import gethostname
import mmcv
import torch
import torch.distributed as dist
......@@ -13,7 +14,11 @@ def get_host_info():
def get_dist_info():
if dist.is_initialized():
if torch.__version__ < '1.0':
initialized = dist._initialized
else:
initialized = dist.is_initialized()
if initialized:
rank = dist.get_rank()
world_size = dist.get_world_size()
else:
......
__version__ = '0.2.2'
__version__ = '0.2.3'
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment