如何在tensorflow 2中使用make_csv_dataset读取多个列作为标签？

import tensorflow as tf
import pandas as pd

CSV_COLUMN_NAMES = ['feature1', 'feature2', 'label1', 'label2']
DEFAULTS = [0, 0, 0, 0]  # 默认值可以根据实际情况进行调整

def parse_csv_row(*row):
    features = dict(zip(CSV_COLUMN_NAMES[:2], row[:2]))  # 将前两列作为特征
    labels = dict(zip(CSV_COLUMN_NAMES[2:], row[2:]))  # 将后两列作为标签
    return features, labels

def load_data(file_pattern, batch_size, shuffle=True):
    dataset = tf.data.experimental.make_csv_dataset(
        file_pattern,
        batch_size=batch_size,
        column_names=CSV_COLUMN_NAMES,
        column_defaults=DEFAULTS,
        label_name=CSV_COLUMN_NAMES[2:],  # 指定标签列名
        select_columns=CSV_COLUMN_NAMES,  # 选择所有列
        header=True,  # CSV文件是否包含标题行
        shuffle=shuffle
    )
    dataset = dataset.map(parse_csv_row)  # 解析CSV行
    return dataset

train_data = load_data('train.csv', batch_size=32)