I think the probability should be sigmoided here https://github.com/IsCoelacanth/TransformingAutoencoder_PyTorch/blob/ed0a2967d2c5ba559fdb2d565b4fa8727b5c9ee3/Capsule.py#L26 Perhaps the generation units should also be sigmoided.