I figured it out. Basically I should use torch.matmul with self.embeds.weight. Doing this solved the saving / loading problem. The hacky way apparently doesn’t work!
logits is a Variable of shape (Batch, Embedding_Dimension) and self.embeds.weight is a Parameter with shape (Vocabulary_Size, Embedding_Dimension)
logits = torch.matmul(logits, self.embeds.weight.t())