KLDiv

class torch_uncertainty.losses.KLDiv(model)

KL divergence loss for Bayesian Neural Networks. Gathers the KL from the modules computed in the forward passes.

Parameters:

model (nn.Module) – Bayesian Neural Network