自定义训练循环中的损失缩放

Question

我目前正在尝试根据tensorflow文档在自定义训练循环中实现损失定标。

    with tf.GradientTape() as tape:
       logits = network(input)  
       loss_value = loss(target, logits)
       loss_value = optimizer.get_scaled_loss(tf.cast(loss_value, tf.float32))

    grads = tape.gradient(loss_value, network.trainable_variables)
    grads = optimizer.get_unscaled_gradients(tf.cast(grads, tf.float32))

但是，get_unscaled_gradients（）方法遇到以下错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [22,250] != values[1].shape = [250] [Op:Pack] name: x

由于我的网络在输入，输出和密集层中具有不同数量的神经元，因此渐变中的梯度张量也具有不同的维度。

如果不支持此功能，如何在自定义循环中实现损失定标？当然，我可以热切地尝试并遍历所有的毕业生，但是，从性能角度来看，这没有任何意义。