Fixed bugs for MNIST