tensorflow如何求模型梯度?

tensorflow框架,想求分类模型model()对输入x的梯度。
在输出层为sigmoid激活函数时无问题,但如果输出层使用softmax函数,该如何求输出对x的梯度?

def grad(x, model):

    x = tf.constant([x], dtype=tf.float32)
    with tf.GradientTape() as tape:
        tape.watch(x)
        y_pred = model(x)
    gradient = tape.gradient(y_pred, x)
    return gradient[0].numpy()

感谢解答。

不知道你这个问题是否已经解决, 如果还没有解决的话:

如果你已经解决了该问题, 非常希望你能够分享一下解决方案, 写成博客, 将相关链接放在评论区, 以帮助更多的人 ^-^