Как я могу включить градиент выходов по входам в функцию потерь, используя Keras?

Question

Мне удалось включить градиент в функцию потерь, но проблема в том, что я не могу выбрать нужный мне c градиент. Модель имеет несколько входов и один выход, и в настоящее время я не уверен, какой из них используется. Модель определяется следующим образом:

input_shape = (5,)
input_tensor = Input(shape = input_shape)
l1 = Dense(400, activation = 'relu')(input_tensor)
l2 = Dense(400, activation = 'relu')(l1)
l3 = Dense(400, activation = 'relu')(l2)
l4 = Dense(400, activation = 'relu')(l3)
output_tensor = Dense(1, activation = 'relu')(l4)
model = Model(input_tensor, output_tensor)


def custom_loss(input_tensor, output_tensor):
    def newloss(y_true, y_pred):
        mse = K.mean(K.square(y_true - y_pred))
        gradients = K.gradients(output_tensor, input_tensor)[0][:,1]
        return mse + K.maximum(-1*gradients, 0)
    return newloss


sgd = keras.optimizers.Adam(lr = 0.001, beta_1 = 0.9, beta_2 = 0.999)
model.compile(loss = custom_loss(input_tensor, output_tensor),
              optimizer = 'sgd',
              metrics = ['mae'])

epochs = 10
batch_size = 100
# Fit the model weights.
history = model.fit(x_train_bs, y_train_bs,
          batch_size=batch_size,
          epochs=epochs,
          verbose=1,
          validation_data=(x_test_bs, y_test_bs))

Исходя из моего понимания, функция K.gradients должна производить 5 производных (из-за 5 входов), но я не могу проиндексировать ее, чтобы выбрать ту, которая мне нужна в частности. Я совершенно неопытен с Керасом, поэтому любая помощь / интуиция будет высоко ценится.