Todo: [] Rewrite the forward pass [] Add the decoder layers [] Strip the weights of the current pretrained weights [] Test it out