High level pyro models- and training loops

Does there exist, or have anyone had any success with building robust, standardized training loops that works like e.g. pytorch-lightning? Ive tried to use pytorch-lightning to train a pyro model with varying sucess. See

It kinda works, but especially the optimizer and the checkpointing is problematic. And I havent tried gpu or multi-gpu yet.

have anyone else tried and got any good success with standardizing this?