如何将train_date.take(1)与Tensorflow一起使用

import tensorflow as tf

# 创建一个简单的数据集
data = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
dataset = tf.data.Dataset.from_tensor_slices(data)

# 创建一个迭代器
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()

# 使用 take 操作取出第一个元素
first_element = next_element.take(1)

# 运行会话查看结果
with tf.Session() as sess:
    try:
        while True:
            print(sess.run(first_element))
    except tf.errors.OutOfRangeError:
        pass

# 确保数据集不为空
if data.shape[0] > 0:
    first_element = next_element.take(1)
else:
    print("数据集为空")

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    dataset = tf.data.Dataset.from_tensor_slices(data).batch(2)
    iterator = dataset.make_initializable_iterator()
    next_element = iterator.get_next()

    # 在每个设备上分别进行 take 操作
    first_element = next_element.take(1)

with tf.Session() as sess:
    sess.run(iterator.initializer)
    try:
        while True:
            print(sess.run(first_element))
    except tf.errors.OutOfRangeError:
        pass