Is it correct to replace a class after self.post_init() function?
I am using BartModel. However, My project is using a 3-dimensional attention mask that is not accepted by BartModel (I already tried to input my attention mask with the size of [batch, seq_length, seq_length] and it pulls out errors) andposition_idsthat is not an argument when using BartModel. Therefore, I add some codes in_expand_mask()function, fix something inBartLearnedPositionalEmbedding()class, andBartModelclass to be suitable for my project. BecauseBartEncoderandBartDecoderclass have_expand_maskfunction andBartLearnedPositionalEmbedding()class, I overrided them to replaceself.embed_positionsinstance with my class that I have already fixed as below:class My_BartEncoder(BartEncoder):
    def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):
        super().__init__(config, embed_tokens)
        embed_dim = config.d_model
        self.embed_positions = BartLearnedPositionalEmbedding(
            config.max_position_embeddings, 
            embed_dim
        )and also addposition_idsargument toforwardfunction because I need it for my newBartLearnedPositionalEmbedding(I did not change the name of the class, just modify the code). And do the same thing withBartDecoderclass. Finally, I changeself.encoderandself.decoderinBartModelclass withMy_BartEncoderandMy_BartDecoderas I do above.I want to ask whether I am doing right, I read thatself.post_init()is executed at the end. However, I inheritedBartModelthen replace 2 instances like thatclass My_BartModel(BartModel):
    def __init__(self, config: BartConfig):
        super().__init__(config)
        self.encoder = My_BartEncoder(config, self.shared)
        self.decoder = My_BartDecoder(config, self.shared)is it legal? And second, I use another names that are notBartEncoderandBartDecoder, so when I initialize model by usingfrom_pretrained, is there any errors or the weights can not recognize the class to put the weights into it? I am a newbie and grateful when you help me.