Nan in loss function

[hieugiaosu]

Thanks for your great repository. But when i try to train with your base pipeline and model, the loss start to produce NaN.