probly.train.bayesian.torch

Collection of torch Bayesian training functions.

Classes

ELBOLoss

Evidence lower bound loss based on [BCKW15].

Functions

collect_kl_divergence

Collect the KL divergence of the Bayesian model by summing the KL divergence of each Bayesian layer.