Input data is in the shape (batch_size, max_len) and output is in the shape (batch_size, max_len, vocab_size)

I am using Adam optimizer with lr of 0.001
I tried training the model taking the mean loss instead of the sum, I am still getting a spike in the loss.