如何在tensorflow中省略4维张量中的零？

import tensorflow as tf

# 输入的4维张量
input_tensor = ...

# 将4维张量展平为2维
flatten_tensor = tf.reshape(input_tensor, [batch_size * height * width, channels])

# 计算每个样本的和
sum_tensor = tf.reduce_sum(flatten_tensor, axis=1)

# 找出非零元素的索引
non_zero_indices = tf.where(tf.not_equal(sum_tensor, 0))

# 根据非零索引重新构建张量
output_tensor = tf.gather(flatten_tensor, non_zero_indices)

# 将张量恢复为原始形状
output_tensor = tf.reshape(output_tensor, [batch_size, height, width, -1])