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

[ghost]

Make train any variant.

[rmsc]

In data_utils.py remove/comment change the Normalize() lines to:

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

[rmsc]

Turns out you should still normalize the data, just edited my answer above.