RuntimeError: Function 'MulBackward0' returned nan values in its 0th output

[albanD]

@tianq01

learning rate and optimizer setting are also ok. As mentioned, the training with the same config using tensorflow implementation does not NaN.

You should check the optimizers doc in details as there might be differences! Momentum in particular IIRC