在tensorflow中，如何计算从csv生成的批的每一列的平均值，该csv在多列中具有NaNs？

import tensorflow as tf

# 定义CSV文件的路径和列名
csv_file = "path/to/csv/file.csv"
column_names = ['col1', 'col2', 'col3']

# 定义解析CSV文件的函数
def parse_csv_line(line):
    # 解析每一行数据
    parsed_line = tf.io.decode_csv(line, record_defaults=[[0.0]] * len(column_names))
    # 返回一个字典，键为列名，值为解析后的数据
    return dict(zip(column_names, parsed_line))

# 读取CSV文件并解析数据
dataset = tf.data.TextLineDataset(csv_file).skip(1)  # 跳过文件头
dataset = dataset.map(parse_csv_line)

# 计算每一列的平均值
column_sums = None
column_counts = None

for batch in dataset.batch(32):
    # 过滤掉包含NaN值的列
    batch = {k: tf.where(tf.math.is_nan(v), tf.zeros_like(v), v) for k, v in batch.items()}
    
    if column_sums is None:
        # 初始化列的累加和和计数
        column_sums = {k: tf.reduce_sum(v) for k, v in batch.items()}
        column_counts = {k: tf.reduce_sum(tf.ones_like(v)) for k, v in batch.items()}
    else:
        # 更新列的累加和和计数
        column_sums = {k: column_sums[k] + tf.reduce_sum(v) for k, v in batch.items()}
        column_counts = {k: column_counts[k] + tf.reduce_sum(tf.ones_like(v)) for k, v in batch.items()}

# 计算每一列的平均值
column_means = {k: column_sums[k] / column_counts[k] for k in column_sums.keys()}

# 打印每一列的平均值
for k, v in column_means.items():
    print(f"Average of {k}: {v}")