Tensorflow while_loop для обучения

В моей задаче мне нужно запустить GD с 1 примером из данных на каждом этапе обучения. Известно, что у session.run () есть издержки, и поэтому тренировать модель слишком долго. В попытке избежать накладных расходов я попытался использовать while_loop и train train для всех данных одним вызовом run (). Но этот подход не работает, и train_op не выполняет даже те из них. Ниже простой пример того, что я делаю:

data = [k*1. for k in range(10)]
tf.reset_default_graph()

i = tf.Variable(0, name='loop_i')
q_x = tf.FIFOQueue(100000, tf.float32)
q_y = tf.FIFOQueue(100000, tf.float32)

x = q_x.dequeue()
y = q_y.dequeue()
w = tf.Variable(0.)
b = tf.Variable(0.)
loss = (tf.add(tf.mul(x, w), b) - y)**2

gs = tf.Variable(0)

train_op = tf.train.GradientDescentOptimizer(0.05).minimize(loss, global_step=gs)

s = tf.Session()
s.run(tf.initialize_all_variables())

def cond(i):
    return i < 10

def body(i):
    return tf.tuple([tf.add(i, 1)], control_inputs=[train_op])


loop = tf.while_loop(cond, body, [i])

for _ in range(1):
    s.run(q_x.enqueue_many((data, )))
    s.run(q_y.enqueue_many((data, )))

s.run(loop)
s.close()

Что я делаю не так? Или есть другое решение этой проблемы со слишком дорогими накладными расходами?

Спасибо!