Finalize the documentation and some refactoring

[constantinpape]

I have updated most doc strings in #463. But the doc strings for the following are still missing:

• torch_em.model.vim: I am not sure how stable this is or if we should update it to newer (vision) mamba implementations.
• torch_em.model.probabilistic_unet: I am also not sure how finalized this is, given that we don't have support for 3D yet.
• torch_em.self_training.loss.ProbabilisticUNetLoss: Ths should be refactored, as this is more general than self-training.
• torch_em.self_training.probabilistic_unet_trainer.py: Should also be refactored.