def create_mems_tf(bsz_per_core):
    mems = [tf.placeholder(dtype=tf.float32, shape=[FLAGS.mem_len, bsz_per_core, FLAGS.d_model]) for layer in range(FLAGS.n_layer)]
    return mems
def create_mems_tf(bsz_per_core):
    mems = [tf.placeholder(dtype=tf.float32, shape=[FLAGS.mem_len, bsz_per_core, FLAGS.d_model]) for layer in range(FLAGS.n_layer)]
    return mems