Update llmfoundry/models/mpt/modeling_mpt.py

Co-authored-by: Sasha Doubov <[email protected]>
mosaicml · Nov 11, 2023 · fea8f95 · fea8f95
1 parent e740386
commit fea8f95
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -618,7 +618,9 @@ def set_input_embeddings(
 
     def get_output_embeddings(
             self) -> Union[SharedEmbedding, nn.Embedding, nn.Linear]:
-        return self.lm_head or self.transformer.get_input_embeddings()
+        if self.lm_head is not None:
+            return self.lm_head
+        return self.transformer.get_input_embeddings()
 
     def set_output_embeddings(
         self, new_embeddings: Union[SharedEmbedding, nn.Embedding,