Now the loss is not divided by the VOCAB_SIZE so it is expected to be higher. If there is a problem with the training, then you might need to reduce the learning_rate. The reason is given that the loss is scaled, then the gradients are scaled as well, so the update steps are larger. So you can lower the learning_rate to avoid that.