Make the mini-batches match for inputs and output-

There seems to be an error in the choice of batches in the train procedure, starting from line 110. Read the comments.