Tensorflow не может продолжать работать, TF2.x

Question

Я использую TF2.x с включенным по умолчанию нетерпеливым исполнением. Однако при использовании пользовательской функции потерь он сообщает, что нетерпеливое выполнение имеет значение Ложь.

tf import veryfing eager execution равно True:

import tensorflow as tf
print(tf.executing_eagerly())

Что печатает True. Я использую дифференцируемое ранжирование из https://github.com/google-research/fast-soft-sort

Я продемонстрирую с набором данных MNIST:

import tensorflow as tf
print(tf.executing_eagerly())
from tf_ops import soft_rank

def rank_loss3(orig, pred):
    orig_shape = tf.shape(orig)
    pred_shape = tf.shape(pred)
    orig_rank = soft_rank(tf.reshape(orig, [orig_shape[1], orig_shape[2]]), regularization_strength=1.0)
    pred_rank = soft_rank(tf.reshape(pred, [pred_shape[1], pred_shape[2]]), regularization_strength=1.0)
    rank_mae = tf.keras.losses.MAE(orig_rank, pred_rank)
    return rank_mae

# https://blog.keras.io/building-autoencoders-in-keras.html
from keras.datasets import mnist
import numpy as np
(x_train, _), (x_test, _) = mnist.load_data()
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
x_train = np.reshape(x_train, (len(x_train), 28, 28, 1))  # adapt this if using `channels_first` image data format
x_test = np.reshape(x_test, (len(x_test), 28, 28, 1))  # adapt this if using `channels_first` image data format

# makes dataset smaller for faster testing
x_train = x_train[0:1000]
x_test = x_test[0:100]

from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D
from keras.models import Model
from keras import backend as K

input_img = Input(shape=(28, 28, 1))  # adapt this if using `channels_first` image data format

x = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)
x = MaxPooling2D((2, 2), padding='same')(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
x = MaxPooling2D((2, 2), padding='same')(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
encoded = MaxPooling2D((2, 2), padding='same')(x)

# at this point the representation is (4, 4, 8) i.e. 128-dimensional

x = Conv2D(8, (3, 3), activation='relu', padding='same')(encoded)
x = UpSampling2D((2, 2))(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
x = UpSampling2D((2, 2))(x)
x = Conv2D(16, (3, 3), activation='relu')(x)
x = UpSampling2D((2, 2))(x)
decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)

autoencoder = Model(input_img, decoded)
# autoencoder.compile(optimizer='adadelta', loss='binary_crossentropy')
autoencoder.compile(optimizer='adadelta', loss=rank_loss3)

autoencoder.fit(x_train, x_train,
                epochs=10,
                batch_size=128,
                shuffle=True,
                validation_data=(x_test, x_test))

Затем я получаю следующую ошибку с указанием этого нетерпеливого выполнение неверно.

Traceback (most recent call last):
  File "small_test.py", line 53, in <module>
    autoencoder.compile(optimizer='adadelta', loss=rank_loss3)
  File "/home/ac/.local/lib/python3.8/site-packages/keras/backend/tensorflow_backend.py", line 75, in symbolic_fn_wrapper
    return func(*args, **kwargs)
  File "/home/ac/.local/lib/python3.8/site-packages/keras/engine/training.py", line 229, in compile
    self.total_loss = self._prepare_total_loss(masks)
  File "/home/ac/.local/lib/python3.8/site-packages/keras/engine/training.py", line 691, in _prepare_total_loss
    output_loss = loss_fn(
  File "/home/ac/.local/lib/python3.8/site-packages/keras/losses.py", line 71, in __call__
    losses = self.call(y_true, y_pred)
  File "/home/ac/.local/lib/python3.8/site-packages/keras/losses.py", line 132, in call
    return self.fn(y_true, y_pred, **self._fn_kwargs)
  File "small_test.py", line 10, in rank_loss3
    orig_rank = soft_rank(tf.reshape(orig, [orig_shape[1], orig_shape[2]]), regularization_strength=1.0)
  File "/work/ecl/fast_soft_sort/tf_ops.py", line 73, in soft_rank
    assert tf.executing_eagerly()
AssertionError

Я не могу вспомнить, где я это читал, но считаю, что когда-то активное выполнение не может быть отключено. Что делает это утверждение ошибочным?