simple-bnn-example.py

import torch
from torch import nn
from blitz.modules import BayesianLinear

class BayesianRegressor(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)
        self.blinear1 = BayesianLinear(input_dim, 64)
        self.blinear2 = BayesianLinear(64, output_dim)
        
    def forward(self, x):
        x_ = self.linear(x)
        x_ = self.blinear1(x_)
        return self.blinear2(x_)