The lack of colour channels in MNIST (and MNIST-like) datasets means we get errors. It looks like the network implementations aren't set up to handle single channel inputs? Could you parameterise them so that they can be used for MNIST? It would be ideal to just write:
# This:
model = PiModel(channels=1)
# Or this
model = PiModel(shape=(28,28,1))
The lack of colour channels in MNIST (and MNIST-like) datasets means we get errors. It looks like the network implementations aren't set up to handle single channel inputs? Could you parameterise them so that they can be used for MNIST? It would be ideal to just write: