NotFoundError: [_Derived_]No gradient defined for op: Einsum on Tensorflow 1.15

I’m using Tensorflow 1.15.2 for making a WSD system, made with BERt in the Embeddings Layer. This is the code that I use for the model

input_word_ids = tf.keras.layers.Input(shape=(64,), dtype=tf.int32, name="input_word_ids")
input_mask = tf.keras.layers.Input(shape=(64,), dtype=tf.int32, name="input_mask")
segment_ids = tf.keras.layers.Input(shape=(64,), dtype=tf.int32, name="segment_ids")
# BERt = BERtLayer()([input_word_ids, input_mask, segment_ids])
bert = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_multi_cased_L-12_H-768_A-12/1", trainable=True)
pooled_output, sequence_output = bert([input_word_ids, input_mask, segment_ids])
# self.vocab_file = bert.resolved_object.vocab_file.asset_path.numpy()
# self.do_lower_case = bert.resolved_object.do_lower_case.numpy()
LSTM = tf.keras.layers.Bidirectional(
    tf.keras.layers.LSTM(
        units=hidden_size,
        dropout=dropout,
        recurrent_dropout=recurrent_dropout,
        return_sequences=True,
        return_state=True
    )
)(sequence_output)

Keras, though, on training, raises the following exception. How can I fix this?

NotFoundError: [_Derived_]No gradient defined for op: Einsum
     [[{{node Func/_36}}]]
     [[training/SGD/gradients/gradients/keras_layer/cond/StatefulPartitionedCall_grad/PartitionedCall/gradients/StatefulPartitionedCall_grad/PartitionedCall/gradients/StatefulPartitionedCall_grad/SymbolicGradient]]

[hammadrizwan]

Yes setting trainable=True triggers the error.

[RobRomijnders]

[I had a similar issue. Not sure if my solution applies. ]

I realized my hub module has several tags. Trying a more appropriate tag did have the gradient op.

Here’s a code sample that you could quickly try: For example, my module has tag ‘“A”’ and tag ‘“B”’. For tag “B”, the gradient op is defined, then the following should work:

with tf.GradientTape() as tape:
  some_input = tf.random.uniform(shape=(13, 224, 224, 3))

  model = hub.KerasLayer("/path/to/hub", tags='B')

  some_output = model(some_input)
  loss = tf.reduce_mean((some_output)**2)

grad = tape.gradient(loss, im)
print(grad)

Running the same with tag ‘“A”’ will fail with the same error.