error: output with shape [1, 28, 28] doesn't match the broadcast shape [3, 28, 28]

[MonroeHe]

I got the error: output with shape [1, 28, 28] doesn't match the broadcast shape [3, 28, 28]
because of this code:
for batch_idx, (x, _) in enumerate(train_loader)

we should change the transform code for mnist dataset:
ERROR
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])

Right
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])