Вход декодера для трансформатора

Я пытаюсь создать вариационный автоэнкодер, используя Longformer от Huggingface. Но я изо всех сил пытаюсь подключить линейный слой к декодеру. В настоящее время подаю заявку

self.decoder(inputs_embeds=x, warning_mask=mask1, global_attention_mask=mask2), тогда как x=encoder.last_hidden_state с x.shape = (batch, sequence_length, hidden_size)

Тренируется, но потери не уменьшаются. Я также применил weight_decay, lr-scheduler и увеличение/уменьшение сложности модели.