Problem! Bayesian NNs classification using Minist dataset

Hello! I got trouble with my Neural Network. This is my first time to create a neural network.
I am trying to use Pyro to solve a classification problem, which is a future topic of Autonomous driving.
And I am trying to use Minist dataset (hello world! in nn) for my first attempt. I use the classification adapted from the tutorial’s regression.

here is the code:

train_loader = torch.utils.data.DataLoader(
datasets.MNIST(‘mnist-data/’, train=True, download=True,
transform=transforms.Compose([transforms.ToTensor(), ])),
batch_size=100, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST(‘mnist-data/’, train=False, transform=transforms.Compose([transforms.ToTensor(), ])
),
batch_size=100, shuffle=True)

class BnnModel(PyroModule):

def __init__(self, input_size=28*28, hidden_size=1024, output_size=10):
    super().__init__()
    self.fc1 = PyroModule[nn.Linear](input_size, hidden_size)
    self.fc1.weight = PyroSample(dist.Normal(0., 1.).expand([hidden_size, input_size]).to_event(2))
    self.fc1.bias = PyroSample(dist.Normal(0., 3.).expand([hidden_size]).to_event(1))
    self.fc2 = PyroModule[nn.Linear](hidden_size, output_size)
    self.fc2.weight = PyroSample(dist.Normal(0., 1.).expand([output_size, hidden_size]).to_event(2))
    self.fc2.bias = PyroSample(dist.Normal(0., 3.).expand([output_size]).to_event(1))
    self.relu = nn.ReLU()

def forward(self, x, y=None):
    output = self.fc1(x)
    output = self.relu(output)
    mean = self.fc2(output).squeeze()
    sigma = pyro.sample("sigma", dist.Uniform(0., 3.))
    with pyro.plate("data"):
        obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
    return mean

hook

net = BnnModel()
auto_guide = pyro.infer.autoguide.AutoNormal(net)
adam = pyro.optim.Adam({“lr”: 0.01})
svi = SVI(net, auto_guide, adam, loss=Trace_ELBO())
pyro.clear_param_store()

num_iterations = 5

for j in range(num_iterations):
loss = 0
count = 0
for batch_id, data in enumerate(train_loader):
ininin, labellabel = data
count += 1
# calculate the loss and take a gradient step
loss = svi.step(ininin.view(-1, 28*28), labellabel)
print("Epoch ", j, " Loss ", loss, “counter =”, count)

However, there are errors:

Traceback (most recent call last):
File “D:\Study app\code\study_pyro\uncertainty1.py”, line 59, in
loss = svi.step(x_train, y_train)
ValueError: Shape mismatch inside plate(‘data’) at site obs dim -1, 20 vs 1000
Trace Shapes:
Param Sites:
Sample Sites:
fc1.weight dist | 20 1
value | 20 1
fc1.bias dist | 20
value | 20
fc2.weight dist | 20 20
value | 20 20
fc2.bias dist | 20
value | 20
fc3.weight dist | 1 20
value | 1 20
fc3.bias dist | 1
value | 1
sigma dist |
value |
Trace Shapes:
Param Sites:
Sample Sites: