如何在TensorFlow中解决'BiasGrad requires tensor size <= int32 max‘InvalidArgumentError？

import tensorflow as tf

# 设置数据类型为tf.float64
tf.keras.backend.set_floatx('float64')

# 构建模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 定义优化器
optimizer = tf.keras.optimizers.Adam()

# 定义损失函数
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

# 定义训练步骤
@tf.function
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        logits = model(inputs)
        loss_value = loss_fn(labels, logits)
    grads = tape.gradient(loss_value, model.trainable_variables)
    grads = [tf.clip_by_value(grad, -2**31, 2**31 - 1) for grad in grads]  # 对梯度进行裁剪
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss_value

# 进行训练
for inputs, labels in train_dataset:
    loss_value = train_step(inputs, labels)
    # 其他训练过程...