VAE KL divergence loss decreases really fast

I am new to training VAEs and I am using it on some 16x16 images, that contains some images from a physics experiment with one or 2 events i.e. the images are mainly black, except for one or 2 regions where we have a Gaussian-like signal around several pixels. Below are the main components of my code:

class View(nn.Module):
def __init__(self, size):
super(View, self).__init__()
self.size = size

def forward(self, tensor):
return tensor.view(self.size)

class Encoder(nn.Module):
def __init__(self, lat_dim):
super().__init__()

N = 32

self.encoder = nn.Sequential(
nn.Conv2d(1, N, 4, 2, 1),          # B, N, 8, 8
nn.ReLU(True),
nn.Conv2d(N, N, 4, 2, 1),          # B,  N, 4, 4
nn.ReLU(True),
nn.Conv2d(N, 2*N, 4, 1),          # B,  N, 4, 4
nn.ReLU(True),
View((-1, 2*N*1*1)),                 # B, 2*N
)

self.mu = nn.Linear(2*N, lat_dim)
self.var = nn.Linear(2*N, lat_dim)

def forward(self, x):

x = self.encoder(x)
z_mu = self.mu(x)
z_var = self.var(x)

return z_mu, z_var

class Decoder(nn.Module):
def __init__(self, lat_dim):
super().__init__()

N = 32

self.decoder = nn.Sequential(
nn.Linear(lat_dim, 2*N),               # B, 256
View((-1, 2*N, 1, 1)),               # B, 256,  1,  1
nn.ReLU(True),
nn.ConvTranspose2d(2*N, N, 4),      # B,  64,  4,  4
nn.ReLU(True),
nn.ConvTranspose2d(N, N, 4, 2, 1), # B,  64,  8,  8
nn.ReLU(True),
nn.ConvTranspose2d(N, 1, 4, 2, 1), # B,  32, 16, 16
)

def forward(self, x):
x = self.decoder(x)

return x

class VAE(nn.Module):
def __init__(self, enc, dec):
super().__init__()

self.enc = enc
self.dec = dec

def forward(self, x):
# encode
z_mu, z_var = self.enc(x)

# sample from the distribution having latent parameters z_mu, z_var
# reparameterize
std = torch.exp(z_var / 2)
eps = torch.randn_like(std)
x_sample = eps*std+z_mu

# decode
predicted = self.dec(x_sample)
return predicted, z_mu, z_var

lat_dim = 10
model_E = Encoder(lat_dim)
model_D = Decoder(lat_dim)
model = VAE(model_E, model_D).cuda()

for epoch in range(801):
model.train()

if epoch%200==0 and epoch>1:
lrs = lrs/10
optimizer = optim.Adam(model.parameters(), lr = lrs, weight_decay=1e-3)

for x, y in zip(train_data, train_data_labels):

x = x.float().cuda()
y = y.float().cuda()

x_sample, z_mu, z_var = model(x)
recon_loss = F.mse_loss(x_sample, x)
kl_loss = torch.mean(-0.5 * torch.sum(1 + z_var - z_mu ** 2 - torch.exp(z_var), dim = 1), dim = 0)

beta = 0.1
loss = recon_loss + beta*kl_loss

loss.backward()
optimizer.step()


The data is normalized to mean 0 and std of 1. When I train it however, the KL divergence loss goes almost instantly (I guess it starts there) to ~10^-5 and keeps going down to around 10^-10, while the reconstruction loss doesn't change from the starting values (around 1). I tried different values for the beta parameter (the weight of the kl loss) and i have to set that to 10^-4 or lower for the reconstruction loss to actually go down. But if I do that the KL loss doesn't change much (which makes sense I guess). But in most papers I read about VAE the beta parameter is between 0 and 1, so I guess I am doing something wrong. Also my data is a lot simpler in both structure and size compared to MNIST (or other complex data sets), so I guess that a VAE code should work. What am I doing wrong? I am using a batch size of 128 and 6000 images for training. The loss mentioned in the post is the training loss. Thank you!