def __init__(self, num_in=784, num_hidden=200, num_out=10, prior_std=1.):
# call to father constructor
super().__init__()
# define prior
prior = Normal(0, prior_std)
# Define layers
# linear layer 1
self.linear_layer = PyroModule[torch.nn.Linear](num_in, num_hidden)
# linear alyer parameters as random variables
self.linear_layer.weights = PyroSample(
prior.expand([num_hidden, num_in]).to_event(2))
self.linear_layer.bias = PyroSample(
prior.expand([num_hidden]).to_event(1))
# linear layer 2
# output dimension is 3 because of the number of classes
self.output_layer = PyroModule[torch.nn.Linear](num_hidden, num_out)
# linear alyer parameters as random variables
self.output_layer.weights = PyroSample(
prior.expand([num_out, num_hidden]).to_event(2))
self.output_layer.bias = PyroSample(
prior.expand([num_out]).to_event(1))
