Solved by djlindsay in post #2 
              
              
                Fixed it. Batch  and train function are now as follows: 
def get_batches(arr_x, arr_y, batch_size):      
    # iterate through the arrays
    prv = 0
    for n in range(batch_size, arr_x.shape[0]+1, batch_size):
        x = arr_x[prv:n,:]
        y = arr_y[prv:n,:]
        prv = n
        yield x, …