Parallel_apply assert len(modules) == len(inputs) AssertionError when using Multi-GPU


I am using Pytorch version 0.4.1 with Python 3.6. I am adapting the transformer model for translation from this site ( The model runs without error with single gpu. However, the AssertionError happens when I use DataParallel for the model. How should I deal with this issue? Thanks for your help in advance. The following snippet of the code:

def run_epoch(data_iter, model, loss_compute):
    "Standard Training and Logging Function"
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    for i, batch in enumerate(data_iter):
        out = model.forward(batch.src, batch.trg, batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)
        total_loss += loss
        total_tokens += batch.ntokens.float()
        tokens += batch.ntokens.float()
        if i % 50 == 1:
            elapsed = time.time() - start
            print("Epoch Step: %d Loss: %f Tokens per Sec: %f" % (i, loss / batch.ntokens.float().item(), tokens / elapsed))
            start = time.time()
            tokens = 0
    return total_loss / total_tokens.float()

loss_compute function

class MultiGPULossCompute:
    "A multi-gpu loss compute and train function."
    def __init__(self, generator, criterion, devices, opt=None, chunk_size=5):
        # Send out to different gpus.
        self.generator = generator
        self.criterion = nn.parallel.replicate(criterion, 
        self.opt = opt
        self.devices = devices
        self.chunk_size = chunk_size
    def __call__(self, out, targets, normalize):
        total = 0.0
        generator = nn.parallel.replicate(self.generator, 
        out_scatter = nn.parallel.scatter(out, 
        out_grad = [[] for _ in out_scatter]
        targets = nn.parallel.scatter(targets, 

        # Divide generating into chunks.
        chunk_size = self.chunk_size
        for i in range(0, out_scatter[0].size(1), chunk_size):
            # Predict distributions
            out_column = [[Variable(o[:, i:i+chunk_size].data, 
                                    requires_grad=self.opt is not None)] 
                           for o in out_scatter]
            gen = nn.parallel.parallel_apply(generator, out_column)

            # Compute loss. 
            y = [(g.contiguous().view(-1, g.size(-1)), 
                  t[:, i:i+chunk_size].contiguous().view(-1)) 
                 for g, t in zip(gen, targets)]
            loss = nn.parallel.parallel_apply(self.criterion, y)

            # Sum and normalize loss
            l = nn.parallel.gather(loss, 
            l = l.sum()[0] / normalize.float()
            total +=[0]

            # Backprop loss to output of transformer
            if self.opt is not None:
                for j, l in enumerate(loss):

        # Backprop all loss through transformer.            
        if self.opt is not None:
            out_grad = [Variable(, dim=1)) for og in out_grad]
            o1 = out
            o2 = nn.parallel.gather(out_grad, 
        return total * normalize.float()