In the Bayesian regression example:
class BayesianRegression(PyroModule): def __init__(self, in_features, out_features): super().__init__() self.linear = PyroModule[nn.Linear](in_features, out_features) self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2)) self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))
I am confused why in the definition of the linear layer the order is
(in_features, out_features) but in the weight definition we do
expand([out_features, in_features]). I’ve been trying to create my own simple NN and I’m getting all kinds of errors related to incorrect dimensions, and I can’t seem to figure it out.