It seems that dataparallel need the output to have a shape, but the loss in pytorch since 0.4 is a scalar
and does’t have a shape. I modified the return code to
return predict, loss.unsqueeze()

and it works fine.