[keras] feature_importance: fixes grad method

keras.py

@@ -1125,13 +1125,13 @@ class LBNLayer(tf.keras.layers.Layer):
 
 
 def feature_importance_grad(model, x=None, **kwargs):
+    inp = [tf.Variable(v) for v in x]
     with tf.GradientTape() as tape:
-        inp = [tf.Variable(v) for v in x]
-        pred = model(inp)
-        ix = np.argsort(-pred, axis=-1)[:, 0]
-        loss = tf.gather(pred, ix, batch_dims=1)
+        pred = model(inp, training=False)
+        ix = np.argsort(pred, axis=-1)[:, -1]
+        decision = tf.gather(pred, ix, batch_dims=1)
 
-    gradients = tape.gradient(loss, inp)  # gradients for decision nodes
+    gradients = tape.gradient(decision, inp)  # gradients for decision nodes
     normed_gradients = [_g * _x for (_g, _x) in zip(gradients, x)]  # normed to input values
 
     mean_gradients = np.concatenate(