WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.

When I try to train a very simple and small model with tensorflow and tensorflow_probability, the code is following

def neg_log_likelihood(y_true, y_hat):
return -y_hat.log_prob(y_true)

inputs = tf.keras.layers.Input(shape=(X_train.shape[1], ))
h1 = tf.keras.layers.BatchNormalization()(inputs)
h2 = tf.keras.layers.Dense(4, activation=tf.nn.relu, name="layer_1")(h1)
rate = tf.keras.layers.Dense(1+1, activation=tf.exp, name="layer_2")(h2)
y = tfp.layers.DistributionLambda(lambda t: tfd.Skellam(rate1=t[..., 0:1], rate2=t[..., 1:]), name="prob")(rate)

model = tf.keras.Model(inputs=inputs, outputs=y)
model.compile(tf.keras.optimizers.Adam(learning_rate=0.01), loss=neg_log_likelihood)
history = model.fit(x=X_train, y=y_train, validation_data=(X_val, y_val), epochs=100, batch_size=128, shuffle=True)

When I start training the model, I get a tensorflow warning
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.

I ignore the WARNING and continue to train the model. However I find the model training is very slow. The ETA is almost 2 min per epoch. The shape of the training set is (121562, 24)

So how can I fix the problem? Thanks

tensorflow version--2.8.0
tensorflow_probability version--0.16.0

Start by reproducing the supplied Keras model and training call with TensorFlow 2.8.0 and TensorFlow Probability 0.16.0, using the reported dataset shape. Inspect when the custom-gradient warning appears and measure the epoch time; done means the warning and reported training slowdown are explained and a verified resolution is documented or implemented.