samim:

        self.mu = z[0:self.latent_dim]
        self.log_sigma = z[self.latent_dim:]



It looks like when you’re setting the self.mu and self.log_sigma you are sampling from the first dimension which is not the feature dimension but the batch size since the input tensor is a 1x800. I believe you need to sample from the second dimension here so try changing this to
        self.mu = z[:,0:self.latent_dim]
        self.log_sigma = z[:,self.latent_dim:]

Screen Shot 2021-03-13 at 11.42.00 AM1644×984 90.7 KB