在循环时将中间张量保存在张量流中

Question

我想保存循环时在张量流中计算的中间张量结果。

首先，我的循环主体

def _loop_body(iteration_,weights,current_state, kernel_results):
    # update weights, current_state, kernel_results

    # weight is shape (100, dtype=float32) tensor
    # current_state and next_state is shape ([100,120], dtype=float32) tensor
    # kernel_results is a collections.named tuple

    return [iter_ + 1, weights, next_state, kernel_results] 

第二，我的while循环是这样的

    [_, weights, current_state, kernel_results] = tf.while_loop(
    cond=lambda iter_, *args: iter_ < num_steps,
    body=_loop_body,
    loop_vars=[
        np.int32(0),  # iter_
        weights,
        current_state,
        previous_kernel_results,
    ],
    parallel_iterations=parallel_iterations)

是否有将中间结果保存在循环中？

过去关于将中间值保存在循环中的帖子在循环中更新的值是标量而我的是2D张量时起作用。