Beta-Variational AutoEncoder can't disentangle

I am working on a dummy example with generated heartbeats, and want to first use a VAE to encode the heartbeats and afterwards a simple classifier.

Problem is when i increase the beta above 0.01, the reconstructions become nonsense (see the first image). And when the beta is low i get a normal autoencoder output with no disentanglement (second image).

I believe the problem may be in my KL divergence or VAE loss function, but i can't seem to find it. In my encoder i do the reparameterization as such:

enc = self.encoder(x,batch_size, x_lenghts)
mu = self.enc2mean(enc)
logv = self.enc2logv(enc)
std = torch.exp(0.5*logv)
z = torch.randn([batch_size,1, self.encoder_hidden_sizes[-1] * (int(self.bidirectional)+1)]).to(self.device)
z = z * std + mu

And i define the VAE loss as:

def VAE_loss(x, reconstruction, mu, logvar, batch_size, latent_dim, beta=0):
    mse = F.mse_loss(x, reconstruction)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    KLD /= (batch_size * latent_dim)
    return mse + beta*KLD

Full standalone code to reproduce the results is here.

Any insights are appreciated!