Solved by abyliss in post #2 
              
              
                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, Embeddi…