@valhalla  Thanks, something like:
class BartExp(nn.Module):
    def __init__(self,config):
        super(BartExp,self).__init__()
        self.bart = BartForConditionalGeneration.from_pretrained(config)
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        decoder_input_ids=None,
        decoder_attention_mask=None,
        head_mask=None,
        decoder_head_mask=None,
        encoder_outputs=None,
        past_key_values=None,
        inputs_embeds=None,
        decoder_inputs_embeds=None,
        labels=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        output = self.bart(input_ids=input_ids,attention_mask=attention_mask,
                     labels=labels,decoder_input_ids=decoder_input_ids,encoder_outputs=encoder_outputs)
        return output

model =  BartExp('facebook/bart-base')
and then when I’m inside a training loop:
output = model(input_ids=original_input_ids,attention_mask=original_attention_mask,
              labels=labels,decoder_input_ids=None,encoder_outputs=None)

It does seem to be working… am i doing it right?