使用tf.data.Dataset.map的正确方法是什么？

import tensorflow as tf

# 定义转换函数
def preprocess_data(element):
    # 对每个元素进行处理
    # ...
    return processed_element

# 创建数据集并加载数据
dataset = tf.data.Dataset.from_tensor_slices(data)

# 应用转换函数
dataset = dataset.map(preprocess_data)

# (可选)对数据集进行进一步处理
dataset = dataset.batch(batch_size)
dataset = dataset.shuffle(buffer_size)

# 迭代数据集
for element in dataset:
    # 处理处理后的数据
    # ...