I read an interesting paper on transformer based variational autoencoder (VAE), so I tried to replicate a simpler model using nn.TransformerEncoderLayer and nn.TransformerDecoderLayer .The architecture is given code :-
embedding_size = 256
n_heads = 8
batch_size = 16
latent_dim = 32

class Transformer_VAE(nn.Module):
    
    def __init__(self, head, vocab_size, embedding_size, latent_dim, device = 'cpu', pad_idx = 0, start_idx = 1, end_idx = 2, unk_idx = 3):
        super(Transformer_VAE, self).__init__()
        self.head = head
        self.embedding_size = embedding_size
        self.vocab_size = vocab_size
        self.latent_dim = latent_dim
        self.embed = WordEmbedding(self.vocab_size, self.embedding_size)
        self.postional_encoding = PostionalEncoding(embedding_size, device)
        self.encoder = nn.TransformerEncoderLayer(self.embedding_size, head)
        self.decoder = nn.TransformerDecoderLayer(self.latent_dim, head)
        self.hidden_to_mean = nn.Linear(self.embedding_size, latent_dim)
        self.hidden_to_logvar = nn.Linear(self.embedding_size, latent_dim)
        self.out_linear = nn.Linear(self.embedding_size, vocab_size)
        
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std
        
    def forward(self, x):
        batch_size, maxlen = x.size()[:2]
        src = self.embed(x)
        x = self.postional_encoding(src)
        x = self.encoder(x)
        print(x.size())
        mean, logvar = self.hidden_to_mean(x), self.hidden_to_logvar(x)
        z = self.reparameterize(mean, logvar)
        out = self.decoder(src, z)
        out = self.out_linear(out)
        return mean, logvar, out

I’m getting an AssertionError:, when I run the self.decoder(src, z) line. From what I understood from the documentation the TransformerDecoderLayer takes the target value and the output of the encoder as its inputs. Since this is a VAE, the target will be the same as the input and the encoder outputs have been mapped to the latent dimension. So these two values are input to the decoder layer.