The hyper-parameter "kl_weight: float = 1.0" leads to the futility of switch loss

[Uigi1600]

Thanks for this useful tool. I used it on my scRNA-seq data and learned a lot about velocity model and VAE code from VeloVI.

However, when i read your _module.py code, i find a unreasonable part in "loss" function:

loss = local_loss + self.penalty_scale * (1 - kl_weight) * end_penalty

This part correspond to article methods:

But since the hyper-parameter "kl_weight" was set to 1.0 in this model, the switch loss "end_penalty" is actually weighted to 0, which means there is only elbo loss in VeloVI model but no switch loss.

Is this a bug or i have any misunderstand of your code or article? Please take some time to explain. it seems like the code is inconsistent to your statement in the article.