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.
I think the probability should be sigmoided here
TransformingAutoencoder_PyTorch/Capsule.py
Line 26 in ed0a296
Perhaps the generation units should also be sigmoided.