Removed scan op which significantly improved the speed. On a GPU it was taking about 3 mins to train on MNIST, now it takes 2 seconds.