Thank you very much. That answers my question exactly.
BTW, I don’t know if it’s allowed or forbid. May I ask you to look at another question I post?

  
    
    
    nn.DataParallel: TypeError: expected sequence object with len >= 0 or a single integer 
  
  
    In my forward function: 
def __call__(self, train=True):
    if train:
        predicted = self.forward(...)
        loss = ....
        return loss # return a single value that's fine
    else:
        predicted = self.forward(...)
        return predicted # expected sequence object with len >= 0 or a single integer
        # predicted with shape [16, 1] on each device and I have 4 GPU

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