Обновите Тензор по индексу в tf. while_loop

Question

Хорошо, поэтому я хочу повторить этот keep[count] = i код Pytorch в Tensorflow, это происходит в tf.while_loop, и count и i являются тензорами.

keep = tf.Variable(tf.zeros(tf.size(scores), tf.int64))
count = 0
idx = tf.argsort(scores, axis=0)#scores.sort(0)  # sort in ascending order
idx = idx[-top_k:]
...

def loop_body(idx, keep, count, ...)
    i = idx[-1]  # index of current largest val
    keep = tf.scatter_update(keep, count, i)

tf.while_loop(loop_cond, loop_body, loop_vars)

Это то, что я пытаюсь сделать keep = tf.scatter_update(keep, count, i), это, к сожалению, поднимает AttributeError: 'Tensor' object has no attribute '_lazy_read