Based on https://github.com/monologg/JointBERT 1Linear, 2Linears, 1Transformer decoder, 2Transformer decoders