Error in VAE Architecture

1

Hi. I’ve been trying to write a VAE for a project, but I found that the loss is stabilizing at a pretty high value. I wanted to know if I was just making a simple error in my VAE.
This is the architecture:
class VAE(nn.Module):
def init(self, z_dim):
super(VAE, self).init()

    # Encoder
    self.encode=nn.Sequential(
    nn.Linear(22,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720)
    )
    self.fc_mu = nn.Linear(720, z_dim)
    self.fc_logvar = nn.Linear(720, z_dim)
    # Decoder
    self.decode=nn.Sequential(
    nn.Linear(z_dim, 720),
    nn.ReLU(),
    nn.Linear(720, 720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,22)
    )
    nn.init.kaiming_uniform_(self.fc_mu.weight)
    nn.init.kaiming_uniform_(self.fc_logvar.weight)
    self.encode.apply(weights_init)
    self.decode.apply(weights_init)

Is there any issue that is immediately apparent that could be causing this error? Thanks!