I have implemented Stacked Capsule Auto-Encoder (Kosiorek et al, 2019) in PyTorch. The original implementation by the authors of paper was created with TensorFlow v1 and DeepMind Sonnet. I could not find a good PyTorch implementation, so hope this provides an easier-to-understand alternative.
I mostly kept the architecture of model and hyper-parameters same with the original implementation. However, some parts are refactored for ease of use.
I am kind of new to PyTorch, so there might be some not-so-best practices . Feel free to give any feedback.