Hi @valhalla, double checked but there might be a hole - here is my training_step function (I use lightning) :
    def training_step(self, batch, batch_id):
        """see lighting docs. Need to make sure that first token is not ignored"""

        decoder_inputs = batch["target_ids"][:, :-1].contiguous()
        decoder_labels = batch["target_ids"][:, 1:].clone()
        decoder_labels[batch["target_ids"][:, 1:] == self.tokenizer.pad_token_id] = -100

        loss = self(source_ids=batch["source_ids"],
                    padding_mask=batch["padding_mask"],
                    decoder_inputs=decoder_inputs,
                    decoder_labels=decoder_labels)[0]

        return {"loss": loss}