VAE model for generating 3d model in pytorch3d but failed.

[univ-esuty]

I'm trying to make VAE model for generating 3d model in pytorch3d but it can't generate high quality 3d mesh. Inputs and outputs of the model are image so simply, I used L2_loss and KL divergence for training. Please let me know if you have any good ideas or find my mistakes.

class VarAutoEncoder(nn.Module):
    def __init__(self, z_dim, num_views, render_s):
        super(VarAutoEncoder, self).__init__()
        
        self.num_views = num_views
        self.render_s = render_s
        
        self.enc = nn.Sequential(
            nn.Conv2d(num_views, 64, 3, 1, padding='same'),
            nn.LeakyReLU(0.3),
            nn.Dropout2d(0.25),
            
            nn.Conv2d(64, 128, 3, 2),
            nn.LeakyReLU(0.3),
            nn.Dropout2d(0.25),
            
            nn.Conv2d(128, 128, 3, 2),
            nn.LeakyReLU(0.3),
            nn.Dropout2d(0.25),
            
            nn.Conv2d(128, 128, 3, 2),
            nn.LeakyReLU(0.3),
            nn.Dropout2d(0.25),
            
            nn.Conv2d(128, 64, 3, 2),
            nn.LeakyReLU(0.3),
            nn.Dropout2d(0.25),
            
            nn.Flatten()
        )
        
        self.encmean = nn.Linear(64*3*3, z_dim)
        self.encvar = nn.Linear(64*3*3, z_dim)
        
        self.dec = nn.Sequential(
            nn.Linear(z_dim, 64*3),
            nn.LeakyReLU(0.3),
            nn.Dropout(0.25),
            
            nn.Linear(64*3, 64*3),
            nn.LeakyReLU(0.3),
            nn.Dropout(0.25),
            
            nn.Linear(64*3, 64*3),
            nn.LeakyReLU(0.3),
            nn.Dropout(0.25),
            
            nn.Linear(64*3, 64*3),
            nn.LeakyReLU(0.3),
            nn.Dropout(0.25),
            
            nn.Linear(64*3, 162*3),
            nn.Sigmoid()
        )
        
    def _encoder(self, x):
        x = self.enc(x)
        mean = self.encmean(x)
        var = F.softplus(self.encvar(x))
        return mean, var
    
    def _sample_z(self, mean, var):
        epsilon = torch.randn(mean.shape).to(device)
        return mean + torch.sqrt(var) * epsilon
 
    def _decoder(self, z):
        new_verts = self.dec(z)
        new_verts = torch.reshape(new_verts, (new_verts.shape[1]//3, 3))
                
        mesh = ico_sphere(2, device)
        verts_size = mesh.verts_packed().shape[0]
        mesh._verts_packed = new_verts
        return self.render_s.render_silhouette(mesh)

    def forward(self, x):
        mean, var = self._encoder(x)
        z = self._sample_z(mean, var)
        x = self._decoder(z)
        return x  # return x, z
    
    def loss(self, x):
        mean, var = self._encoder(x)
        KL = -0.5 * torch.mean(torch.sum(1 + torch.log(var) - mean**2 - var))
        z = self._sample_z(mean, var)
        y = self._decoder(z)
        
        reconstruction = torch.mean(torch.sum(x * torch.log(y) + (1 - x) * torch.log(1 - y)))
        lower_bound = [-KL, reconstruction]                                      
        return -sum(lower_bound)

You can check the full code here link