Предварительная загрузка следующей мини-партии массива NumPy во время обучения (чтобы получить 100% использование графического процессора) с помощью tf.Data

Question

У меня есть огромная матрица NumPy, уже загруженная в память. Я хотел бы уместить (выполнить градиентный спуск) 20 000 строк одновременно, но при этом следующие 20 000 строк будут скопированы в память GPU, поэтому GPU будет постоянно снабжаться данными, достигая 100% использования GPU.

Я пытался сделать это с новым tf.Data API. Я написал код ниже, но загрузка едва достигает 50%. Я не совсем уверен, связано ли это с моей ошибкой или это ошибка в Tensorflow.

x = tf.placeholder(shape=[None,xSize], dtype=typ)
y = tf.placeholder(shape=[None,ySize], dtype=typ)

# xT and yT are huge NumPy arrays
def trainOneEpoch(xT, yT):
    batchSize = 20000
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.batch(batchSize).prefetch(2)
    iterator = dataset.make_initializable_iterator()
    sess.run(iterator.initializer, {x: xT, y: yT})

    while True:
        try:
            xS, yS = sess.run(iterator.get_next())
            sess.run(trainer, {x: xS, y: yS})
        except tf.errors.OutOfRangeError:
            break

Как заставить это работать (с tf.Data или без)?