using self.shift_labels instead of self.model.transformer.shift_label…

…s in loss (#1211) Co-authored-by: Daniel King <[email protected]>
mosaicml · May 15, 2024 · 8cd23d5 · 8cd23d5
1 parent cfee4e4
commit 8cd23d5
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
@@ -1157,7 +1157,7 @@ def forward(self, batch: MutableMapping) -> CausalLMOutputWithPast:
 
     def loss(self, outputs: CausalLMOutputWithPast,
              batch: Mapping) -> Union[dict, torch.Tensor]:
-        if self.model.transformer.shift_labels:
+        if self.shift_labels:
             targets = self.get_targets(batch)
         else:
             targets = batch['labels']