diff --git a/dataset.py b/dataset.py
index b6ffd0b3..e6b10f6a 100644
--- a/dataset.py
+++ b/dataset.py
@@ -1,14 +1,21 @@
 import os
 import json
 import torch
-import lib.utils.data as torchdata
+import math
+import random
+# import lib.utils.data as torchdata
 import cv2
 from torchvision import transforms
 import numpy as np
 from scipy.misc import imresize
 
 
-class BaseDataset(torchdata.Dataset):
+def user_collate_fn(batch):
+    assert(len(batch) == 1)
+    return batch[0]
+
+
+class BaseDataset(torch.utils.data.Dataset):
     def __init__(self, odgt, opt, **kwargs):
         # parse options
         self.imgSizes = opt.imgSizes
@@ -24,20 +31,24 @@ def __init__(self, odgt, opt, **kwargs):
             mean=[102.9801, 115.9465, 122.7717],
             std=[1., 1., 1.])
 
-    def parse_input_list(self, odgt, max_sample=-1, start_idx=-1, end_idx=-1):
+    def parse_input_list(self, odgt, world_size=1, rank=0, start_idx=-1, end_idx=-1):
         if isinstance(odgt, list):
             self.list_sample = odgt
         elif isinstance(odgt, str):
             self.list_sample = [json.loads(x.rstrip()) for x in open(odgt, 'r')]
 
-        if max_sample > 0:
-            self.list_sample = self.list_sample[0:max_sample]
-        if start_idx >= 0 and end_idx >= 0:     # divide file list
-            self.list_sample = self.list_sample[start_idx:end_idx]
+        num_total = len(self.list_sample)
+        if world_size > 1:
+            self.num_sample = int(math.ceil(num_total * 1.0 / world_size))
+            self.start_idx = rank * self.num_sample
+            self.end_idx = min(self.start_idx + self.num_sample, num_total)
+        else:
+            self.start_idx = 0
+            self.end_idx = num_total
 
-        self.num_sample = len(self.list_sample)
-        assert self.num_sample > 0
-        print('# samples: {}'.format(self.num_sample))
+        # assert self.num_sample > 0
+        print('Dataset Samples #total: {}, #process [{}]: {}-{}'
+              .format(num_total, rank, self.start_idx, self.end_idx))
 
     def img_transform(self, img):
         # image to float
@@ -64,8 +75,11 @@ def __init__(self, root_dataset, odgt, opt, batch_per_gpu=1, **kwargs):
         self.batch_record_list = [[], []]
 
         # override dataset length when trainig with batch_per_gpu > 1
-        self.cur_idx = 0
-        self.if_shuffled = False
+        self.cur_idx = self.start_idx
+        # self.if_shuffled = False
+
+    def shuffle(self, seed):
+        random.Random(seed).shuffle(self.list_sample)
 
     def _get_sub_batch(self):
         while True:
@@ -78,9 +92,9 @@ def _get_sub_batch(self):
 
             # update current sample pointer
             self.cur_idx += 1
-            if self.cur_idx >= self.num_sample:
-                self.cur_idx = 0
-                np.random.shuffle(self.list_sample)
+            if self.cur_idx >= self.end_idx:
+                self.cur_idx = self.start_idx
+                # np.random.shuffle(self.list_sample)
 
             if len(self.batch_record_list[0]) == self.batch_per_gpu:
                 batch_records = self.batch_record_list[0]
@@ -94,9 +108,9 @@ def _get_sub_batch(self):
 
     def __getitem__(self, index):
         # NOTE: random shuffle for the first time. shuffle in __init__ is useless
-        if not self.if_shuffled:
-            np.random.shuffle(self.list_sample)
-            self.if_shuffled = True
+        # if not self.if_shuffled:
+        #     np.random.shuffle(self.list_sample)
+        #     self.if_shuffled = True
 
         # get sub-batch candidates
         batch_records = self._get_sub_batch()
@@ -173,14 +187,15 @@ def __getitem__(self, index):
             batch_images[i][:, :img.shape[1], :img.shape[2]] = img
             batch_segms[i][:segm.shape[0], :segm.shape[1]] = torch.from_numpy(segm.astype(np.int)).long()
 
-        batch_segms = batch_segms - 1 # label from -1 to 149
+        batch_segms = batch_segms - 1   # label from -1 to 149 for ADE
         output = dict()
         output['img_data'] = batch_images
         output['seg_label'] = batch_segms
         return output
 
     def __len__(self):
-        return int(1e10) # It's a fake length due to the trick that every loader maintains its own list
+        return self.num_sample
+        # return int(1e10) # It's a fake length due to the trick that every loader maintains its own list
         #return self.num_sampleclass
 
 
diff --git a/models/mobilenet.py b/models/mobilenet.py
index 8a6db071..27bf0f96 100644
--- a/models/mobilenet.py
+++ b/models/mobilenet.py
@@ -8,7 +8,6 @@
 import torch
 import torch.nn as nn
 import math
-from lib.nn import SynchronizedBatchNorm2d
 
 try:
     from urllib import urlretrieve
@@ -27,7 +26,7 @@
 def conv_bn(inp, oup, stride):
     return nn.Sequential(
         nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
-        SynchronizedBatchNorm2d(oup),
+        nn.BatchNorm2d(oup),
         nn.ReLU6(inplace=True)
     )
 
@@ -35,7 +34,7 @@ def conv_bn(inp, oup, stride):
 def conv_1x1_bn(inp, oup):
     return nn.Sequential(
         nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
-        SynchronizedBatchNorm2d(oup),
+        nn.BatchNorm2d(oup),
         nn.ReLU6(inplace=True)
     )
 
@@ -53,25 +52,25 @@ def __init__(self, inp, oup, stride, expand_ratio):
             self.conv = nn.Sequential(
                 # dw
                 nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
-                SynchronizedBatchNorm2d(hidden_dim),
+                nn.BatchNorm2d(hidden_dim),
                 nn.ReLU6(inplace=True),
                 # pw-linear
                 nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-                SynchronizedBatchNorm2d(oup),
+                nn.BatchNorm2d(oup),
             )
         else:
             self.conv = nn.Sequential(
                 # pw
                 nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
-                SynchronizedBatchNorm2d(hidden_dim),
+                nn.BatchNorm2d(hidden_dim),
                 nn.ReLU6(inplace=True),
                 # dw
                 nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
-                SynchronizedBatchNorm2d(hidden_dim),
+                nn.BatchNorm2d(hidden_dim),
                 nn.ReLU6(inplace=True),
                 # pw-linear
                 nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-                SynchronizedBatchNorm2d(oup),
+                nn.BatchNorm2d(oup),
             )
 
     def forward(self, x):
@@ -138,7 +137,7 @@ def _initialize_weights(self):
                 m.weight.data.normal_(0, math.sqrt(2. / n))
                 if m.bias is not None:
                     m.bias.data.zero_()
-            elif isinstance(m, SynchronizedBatchNorm2d):
+            elif isinstance(m, nn.BatchNorm2d):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()
             elif isinstance(m, nn.Linear):
@@ -168,4 +167,3 @@ def load_url(url, model_dir='./pretrained', map_location=None):
         sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
         urlretrieve(url, cached_file)
     return torch.load(cached_file, map_location=map_location)
-
diff --git a/models/models.py b/models/models.py
index 1b8723bd..5d6c1037 100644
--- a/models/models.py
+++ b/models/models.py
@@ -2,7 +2,6 @@
 import torch.nn as nn
 import torchvision
 from . import resnet, resnext, mobilenet
-from lib.nn import SynchronizedBatchNorm2d
 
 
 class SegmentationModuleBase(nn.Module):
@@ -56,7 +55,7 @@ def conv3x3(in_planes, out_planes, stride=1, has_bias=False):
 def conv3x3_bn_relu(in_planes, out_planes, stride=1):
     return nn.Sequential(
             conv3x3(in_planes, out_planes, stride),
-            SynchronizedBatchNorm2d(out_planes),
+            nn.BatchNorm2d(out_planes),
             nn.ReLU(inplace=True),
             )
 
@@ -395,7 +394,7 @@ def __init__(self, num_class=150, fc_dim=4096,
             self.ppm.append(nn.Sequential(
                 nn.AdaptiveAvgPool2d(scale),
                 nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
-                SynchronizedBatchNorm2d(512),
+                nn.BatchNorm2d(512),
                 nn.ReLU(inplace=True)
             ))
         self.ppm = nn.ModuleList(self.ppm)
@@ -403,7 +402,7 @@ def __init__(self, num_class=150, fc_dim=4096,
         self.conv_last = nn.Sequential(
             nn.Conv2d(fc_dim+len(pool_scales)*512, 512,
                       kernel_size=3, padding=1, bias=False),
-            SynchronizedBatchNorm2d(512),
+            nn.BatchNorm2d(512),
             nn.ReLU(inplace=True),
             nn.Dropout2d(0.1),
             nn.Conv2d(512, num_class, kernel_size=1)
@@ -444,7 +443,7 @@ def __init__(self, num_class=150, fc_dim=4096,
             self.ppm.append(nn.Sequential(
                 nn.AdaptiveAvgPool2d(scale),
                 nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
-                SynchronizedBatchNorm2d(512),
+                nn.BatchNorm2d(512),
                 nn.ReLU(inplace=True)
             ))
         self.ppm = nn.ModuleList(self.ppm)
@@ -453,7 +452,7 @@ def __init__(self, num_class=150, fc_dim=4096,
         self.conv_last = nn.Sequential(
             nn.Conv2d(fc_dim+len(pool_scales)*512, 512,
                       kernel_size=3, padding=1, bias=False),
-            SynchronizedBatchNorm2d(512),
+            nn.BatchNorm2d(512),
             nn.ReLU(inplace=True),
             nn.Dropout2d(0.1),
             nn.Conv2d(512, num_class, kernel_size=1)
@@ -509,7 +508,7 @@ def __init__(self, num_class=150, fc_dim=4096,
             self.ppm_pooling.append(nn.AdaptiveAvgPool2d(scale))
             self.ppm_conv.append(nn.Sequential(
                 nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
-                SynchronizedBatchNorm2d(512),
+                nn.BatchNorm2d(512),
                 nn.ReLU(inplace=True)
             ))
         self.ppm_pooling = nn.ModuleList(self.ppm_pooling)
@@ -521,7 +520,7 @@ def __init__(self, num_class=150, fc_dim=4096,
         for fpn_inplane in fpn_inplanes[:-1]: # skip the top layer
             self.fpn_in.append(nn.Sequential(
                 nn.Conv2d(fpn_inplane, fpn_dim, kernel_size=1, bias=False),
-                SynchronizedBatchNorm2d(fpn_dim),
+                nn.BatchNorm2d(fpn_dim),
                 nn.ReLU(inplace=True)
             ))
         self.fpn_in = nn.ModuleList(self.fpn_in)
diff --git a/models/resnet.py b/models/resnet.py
index a6674d55..1355738a 100644
--- a/models/resnet.py
+++ b/models/resnet.py
@@ -3,7 +3,6 @@
 import torch
 import torch.nn as nn
 import math
-from lib.nn import SynchronizedBatchNorm2d
 
 try:
     from urllib import urlretrieve
@@ -33,10 +32,10 @@ class BasicBlock(nn.Module):
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super(BasicBlock, self).__init__()
         self.conv1 = conv3x3(inplanes, planes, stride)
-        self.bn1 = SynchronizedBatchNorm2d(planes)
+        self.bn1 = nn.BatchNorm2d(planes)
         self.relu = nn.ReLU(inplace=True)
         self.conv2 = conv3x3(planes, planes)
-        self.bn2 = SynchronizedBatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
         self.downsample = downsample
         self.stride = stride
 
@@ -65,12 +64,12 @@ class Bottleneck(nn.Module):
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super(Bottleneck, self).__init__()
         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = SynchronizedBatchNorm2d(planes)
+        self.bn1 = nn.BatchNorm2d(planes)
         self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                                padding=1, bias=False)
-        self.bn2 = SynchronizedBatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
         self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
-        self.bn3 = SynchronizedBatchNorm2d(planes * 4)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
         self.relu = nn.ReLU(inplace=True)
         self.downsample = downsample
         self.stride = stride
@@ -104,13 +103,13 @@ def __init__(self, block, layers, num_classes=1000):
         self.inplanes = 128
         super(ResNet, self).__init__()
         self.conv1 = conv3x3(3, 64, stride=2)
-        self.bn1 = SynchronizedBatchNorm2d(64)
+        self.bn1 = nn.BatchNorm2d(64)
         self.relu1 = nn.ReLU(inplace=True)
         self.conv2 = conv3x3(64, 64)
-        self.bn2 = SynchronizedBatchNorm2d(64)
+        self.bn2 = nn.BatchNorm2d(64)
         self.relu2 = nn.ReLU(inplace=True)
         self.conv3 = conv3x3(64, 128)
-        self.bn3 = SynchronizedBatchNorm2d(128)
+        self.bn3 = nn.BatchNorm2d(128)
         self.relu3 = nn.ReLU(inplace=True)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
 
@@ -125,7 +124,7 @@ def __init__(self, block, layers, num_classes=1000):
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                 m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, SynchronizedBatchNorm2d):
+            elif isinstance(m, nn.BatchNorm2d):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()
 
@@ -135,7 +134,7 @@ def _make_layer(self, block, planes, blocks, stride=1):
             downsample = nn.Sequential(
                 nn.Conv2d(self.inplanes, planes * block.expansion,
                           kernel_size=1, stride=stride, bias=False),
-                SynchronizedBatchNorm2d(planes * block.expansion),
+                nn.BatchNorm2d(planes * block.expansion),
             )
 
         layers = []
diff --git a/models/resnext.py b/models/resnext.py
index 931ef95f..d74d8999 100644
--- a/models/resnext.py
+++ b/models/resnext.py
@@ -3,7 +3,6 @@
 import torch
 import torch.nn as nn
 import math
-from lib.nn import SynchronizedBatchNorm2d
 
 try:
     from urllib import urlretrieve
@@ -32,12 +31,12 @@ class GroupBottleneck(nn.Module):
     def __init__(self, inplanes, planes, stride=1, groups=1, downsample=None):
         super(GroupBottleneck, self).__init__()
         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = SynchronizedBatchNorm2d(planes)
+        self.bn1 = nn.BatchNorm2d(planes)
         self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                                padding=1, groups=groups, bias=False)
-        self.bn2 = SynchronizedBatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
         self.conv3 = nn.Conv2d(planes, planes * 2, kernel_size=1, bias=False)
-        self.bn3 = SynchronizedBatchNorm2d(planes * 2)
+        self.bn3 = nn.BatchNorm2d(planes * 2)
         self.relu = nn.ReLU(inplace=True)
         self.downsample = downsample
         self.stride = stride
@@ -71,13 +70,13 @@ def __init__(self, block, layers, groups=32, num_classes=1000):
         self.inplanes = 128
         super(ResNeXt, self).__init__()
         self.conv1 = conv3x3(3, 64, stride=2)
-        self.bn1 = SynchronizedBatchNorm2d(64)
+        self.bn1 = nn.BatchNorm2d(64)
         self.relu1 = nn.ReLU(inplace=True)
         self.conv2 = conv3x3(64, 64)
-        self.bn2 = SynchronizedBatchNorm2d(64)
+        self.bn2 = nn.BatchNorm2d(64)
         self.relu2 = nn.ReLU(inplace=True)
         self.conv3 = conv3x3(64, 128)
-        self.bn3 = SynchronizedBatchNorm2d(128)
+        self.bn3 = nn.BatchNorm2d(128)
         self.relu3 = nn.ReLU(inplace=True)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
 
@@ -92,7 +91,7 @@ def __init__(self, block, layers, groups=32, num_classes=1000):
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels // m.groups
                 m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, SynchronizedBatchNorm2d):
+            elif isinstance(m, nn.BatchNorm2d):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()
 
@@ -102,7 +101,7 @@ def _make_layer(self, block, planes, blocks, stride=1, groups=1):
             downsample = nn.Sequential(
                 nn.Conv2d(self.inplanes, planes * block.expansion,
                           kernel_size=1, stride=stride, bias=False),
-                SynchronizedBatchNorm2d(planes * block.expansion),
+                nn.BatchNorm2d(planes * block.expansion),
             )
 
         layers = []
diff --git a/train.py b/train.py
index 4ddc3be1..7ffd096c 100644
--- a/train.py
+++ b/train.py
@@ -8,17 +8,18 @@
 # Numerical libs
 import torch
 import torch.nn as nn
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import torch.utils.data.distributed
 # Our libs
 from config import cfg
-from dataset import TrainDataset
+from dataset import TrainDataset, user_collate_fn
 from models import ModelBuilder, SegmentationModule
 from utils import AverageMeter, parse_devices, setup_logger
-from lib.nn import UserScatteredDataParallel, user_scattered_collate, patch_replication_callback
-import lib.utils.data as torchdata
 
 
 # train one epoch
-def train(segmentation_module, iterator, optimizers, history, epoch, cfg):
+def train(segmentation_module, loader, optimizers, history, epoch, cfg):
     batch_time = AverageMeter()
     data_time = AverageMeter()
     ave_total_loss = AverageMeter()
@@ -28,10 +29,15 @@ def train(segmentation_module, iterator, optimizers, history, epoch, cfg):
 
     # main loop
     tic = time.time()
+    iterator = iter(loader)
     for i in range(cfg.TRAIN.epoch_iters):
+        # set learning rate
+        cur_iter = i + (epoch - 1) * cfg.TRAIN.epoch_iters
+        adjust_learning_rate(optimizers, cur_iter, cfg)
+
+        # load data
         batch_data = next(iterator)
         data_time.update(time.time() - tic)
-
         segmentation_module.zero_grad()
 
         # forward pass
@@ -67,10 +73,6 @@ def train(segmentation_module, iterator, optimizers, history, epoch, cfg):
             history['train']['loss'].append(loss.data.item())
             history['train']['acc'].append(acc.data.item())
 
-        # adjust learning rate
-        cur_iter = i + (epoch - 1) * cfg.TRAIN.epoch_iters
-        adjust_learning_rate(optimizers, cur_iter, cfg)
-
 
 def checkpoint(nets, history, cfg, epoch_num):
     print('Saving checkpoints...')
@@ -84,10 +86,10 @@ def checkpoint(nets, history, cfg, epoch_num):
         '{}/history_epoch_{}.pth'.format(cfg.DIR, epoch_num))
     torch.save(
         dict_encoder,
-       '{}/encoder_epoch_{}.pth'.format(cfg.DIR, epoch_num))
+        '{}/encoder_epoch_{}.pth'.format(cfg.DIR, epoch_num))
     torch.save(
         dict_decoder,
-       '{}/decoder_epoch_{}.pth'.format(cfg.DIR, epoch_num))
+        '{}/decoder_epoch_{}.pth'.format(cfg.DIR, epoch_num))
 
 
 def group_weight(module):
@@ -140,7 +142,15 @@ def adjust_learning_rate(optimizers, cur_iter, cfg):
         param_group['lr'] = cfg.TRAIN.running_lr_decoder
 
 
-def main(cfg, gpus):
+def main_worker(rank, cfg, rank2gpu):
+    # Param setup
+    num_gpus = len(rank2gpu.keys())
+    gpu = rank2gpu[rank]
+    print("Launch GPU: {} for training".format(gpu))
+    dist.init_process_group(
+        backend='nccl', init_method='tcp://127.0.0.1:1234',
+        world_size=num_gpus, rank=rank)
+
     # Network Builders
     builder = ModelBuilder()
     net_encoder = builder.build_encoder(
@@ -161,35 +171,27 @@ def main(cfg, gpus):
     else:
         segmentation_module = SegmentationModule(
             net_encoder, net_decoder, crit)
+    segmentation_module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(segmentation_module)
+
+    torch.cuda.set_device(gpu)
+    segmentation_module.cuda(gpu)
+    segmentation_module = torch.nn.parallel.DistributedDataParallel(segmentation_module, device_ids=[gpu])
 
     # Dataset and Loader
     dataset_train = TrainDataset(
         cfg.DATASET.root_dataset,
         cfg.DATASET.list_train,
         cfg.DATASET,
-        batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
+        batch_per_gpu=cfg.TRAIN.batch_size_per_gpu,
+        world_size=num_gpus, rank=rank)
 
-    loader_train = torchdata.DataLoader(
+    loader_train = torch.utils.data.DataLoader(
         dataset_train,
-        batch_size=len(gpus),  # we have modified data_parallel
-        shuffle=False,  # we do not use this param
-        collate_fn=user_scattered_collate,
-        num_workers=cfg.TRAIN.workers,
-        drop_last=True,
+        batch_size=1,   # modified: each batch is a dict of multiple samples
+        collate_fn=user_collate_fn,
+        num_workers=cfg.TRAIN.workers // num_gpus,
+        drop_last=False,
         pin_memory=True)
-    print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
-
-    # create loader iterator
-    iterator_train = iter(loader_train)
-
-    # load nets into gpu
-    if len(gpus) > 1:
-        segmentation_module = UserScatteredDataParallel(
-            segmentation_module,
-            device_ids=gpus)
-        # For sync bn
-        patch_replication_callback(segmentation_module)
-    segmentation_module.cuda()
 
     # Set up optimizers
     nets = (net_encoder, net_decoder, crit)
@@ -198,11 +200,20 @@ def main(cfg, gpus):
     # Main loop
     history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
 
+    cfg.TRAIN.epoch_iters = len(dataset_train) // cfg.TRAIN.batch_size_per_gpu
+    cfg.TRAIN.max_iters = cfg.TRAIN.epoch_iters * cfg.TRAIN.num_epoch
     for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
-        train(segmentation_module, iterator_train, optimizers, history, epoch+1, cfg)
+        # deterministic data shuffling
+        dataset_train.shuffle(epoch)
+
+        # train one epoch
+        train(segmentation_module, loader_train, optimizers, history, epoch+1, cfg)
 
         # checkpointing
-        checkpoint(nets, history, cfg, epoch+1)
+        if rank == 0:
+            if ((epoch+1) % cfg.TRAIN.save_freq == 0) \
+                    or ((epoch+1) == cfg.TRAIN.num_epoch):
+                checkpoint(nets, history, cfg, epoch+1)
 
     print('Training Done!')
 
@@ -221,6 +232,12 @@ def main(cfg, gpus):
         help="path to config file",
         type=str,
     )
+    parser.add_argument(
+        "--distributed",
+        default=1,
+        type=int,
+        help="using multiprocessing distributed training"
+    )
     parser.add_argument(
         "--gpus",
         default="0-3",
@@ -255,11 +272,9 @@ def main(cfg, gpus):
     num_gpus = len(gpus)
     cfg.TRAIN.batch_size = num_gpus * cfg.TRAIN.batch_size_per_gpu
 
-    cfg.TRAIN.max_iters = cfg.TRAIN.epoch_iters * cfg.TRAIN.num_epoch
-    cfg.TRAIN.running_lr_encoder = cfg.TRAIN.lr_encoder
-    cfg.TRAIN.running_lr_decoder = cfg.TRAIN.lr_decoder
-
     random.seed(cfg.TRAIN.seed)
     torch.manual_seed(cfg.TRAIN.seed)
 
-    main(cfg, gpus)
+    rank2gpu = {i: gpus[i] for i in range(num_gpus)}
+    mp.spawn(main_worker, nprocs=num_gpus, args=(cfg, rank2gpu))
+    # main_worker(cfg, gpus)
