fix bug

mosaicml · Nov 9, 2023 · c5391c3 · c5391c3
1 parent 5fb3e29
commit c5391c3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -666,7 +666,7 @@ def forward(
         )
 
         out = outputs.last_hidden_state.to(self.transformer.wte.weight.device)
-        if self.unembed is not None:
+        if self.transformer.unembed is not None:
             logits = self.transformer.unembed(out)
         else:
             # move outputs to same device as weights for token embedding