there is no such functionality. the closest would be tyxe which offers some pyro-based machinery for bayesian neural networks.
generally speaking we choose to keep most bayesian deep learning techniques out of scope, essentially because they tend to be more ad hoc than not and so it’s difficult to integrate them coherently into a probabilistic programming framework. by contrast it’s clear how e.g. discrete enumeration, stochastic variational inference, and map can be coherently combined. also we have to choose our battles, and it’s likely that packages like torch.optim.swa_utils are better suited for these sorts of things.