Yes you have understood my problem.
I ran the model with torch.set_flush_denormal(True) on the cpu version and the results were the same.
I am using FP32 to train my model.
I am averaging the word embeddings inside one of the functions -
mention_embeddings = self.embeddings(mention_tokens)
mention_lengths = x['mention_tokens_length'].float().to(self.device)

average_mention_embeddings = torch.sum(mention_embeddings, 1) / mention_lengths.unsqueeze(1)

Do you think this can cause any issue?
P.S. I am using word embeddings to train the model so the checkpoint is a little large (1GB ish). Let me know if you still want it.