Replace self.get_targets

mosaicml · Jul 23, 2024 · 6f4f58b · 6f4f58b
1 parent 118e96c
commit 6f4f58b
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -1396,9 +1396,7 @@ def config_class(self) -> Type[MPTConfig]:
         return MPTConfig
 
     def get_targets(self, batch: Mapping) -> torch.Tensor:
-        targets = torch.roll(batch['labels'], shifts=-1)
-        targets[:, -1] = -100
-        return targets
+        return get_targets(batch['labels'])
 
     def forward(self, batch: MutableMapping) -> CausalLMOutputWithPast:
         if self.config.ffn_config['ffn_type'] in ffns_with_megablocks: