nn_kl_divergence.py

@variational_estimator
class BayesianRegressor(nn.Module):
    def __init__(self,):
        super().__init__()
        #self.linear = nn.Linear(input_dim, output_dim)
        self.blinear1 = BayesianLinear(input_dim, 512)
        self.blinear2 = BayesianLinear(512, output_dim)
        
    def forward(self, x):
        x_ = self.blinear1(x)
        return self.blinear2(x_)
      
reg = BayesianRegressor()
preds = reg(dataset)
fit_loss = criterion(preds, labels)
complexity_loss = reg.nn_kl_divergence()

true_loss = fit_loss + complecity_loss