In my forward function:
def __call__(self, train=True):
    if train:
        predicted = self.forward(...)
        loss = ....
        return loss # return a single value that's fine
        # loss.size() = the number of my GPUs.
    else:
        predicted = self.forward(...)
        return predicted # expected sequence object with len >= 0 or a single integer
        # In validation step I want to return the whole predict labels for other purpose
        # predicted with shape [16, 1] on each device and I have 4 GPU

My code works before model=nn.DataParallel(model).