Tensorflow ValueError: веса не могут быть переданы значениям

Question

Я пытаюсь обучить нейронную сеть с Tensorflow (2.1.0) и Keras, но при утверждении model.compile я получаю эту ошибку:

ValueError: weights не может быть передано значения. values.rank = 0. weights.rank = 2. values.shape = (). weights.shape = (None, 20).

Я обучал нейронную сеть на Google Colab и все было в порядке, но теперь я пытаюсь обучить NN на другой удаленной машине и тот же код не работает.

Потеря:

def get_top_one_probability(vector):
  return (K.exp(vector) / K.sum(K.exp(vector)))

def listnet_loss(real_labels, predicted_labels):
  return -K.sum(get_top_one_probability(real_labels) * tf.math.log(get_top_one_probability(predicted_labels)))

Архитектура модели:

model = tf.keras.Sequential()
model.add(Masking(mask_value=0., input_shape=(timesteps, features)))
model.add(Bidirectional(LSTM(units=100, return_sequences=True), input_shape=(timesteps, features)))
model.add(Dense(1, activation='sigmoid'))

model.summary()
model.compile(loss=listnet_loss, optimizer=keras.optimizers.Adam(learning_rate=0.00005, beta_1=0.9, beta_2=0.999, amsgrad=False))

ОБНОВЛЕНИЕ

I обратите внимание, что в Colab я на самом деле использую версию 2.2.0-rc2

Что происходит? Заранее спасибо.