TensorFlow 2 tf.function decorator

Question

У меня есть TensorFlow 2.0 и Python 3.7.5.

Я написал следующий код для выполнения мини-градиентного спуска:

@tf.function
def train_one_step(model, mask_model, optimizer, x, y):
    '''
    Function to compute one step of gradient descent optimization
    '''
    with tf.GradientTape() as tape:
        # Make predictions using defined model-
        y_pred = model(x)

        # Compute loss-
        loss = loss_fn(y, y_pred)

    # Compute gradients wrt defined loss and weights and biases-
    grads = tape.gradient(loss, model.trainable_variables)

    # type(grads)
    # list

    # List to hold element-wise multiplication between-
    # computed gradient and masks-
    grad_mask_mul = []

    # Perform element-wise multiplication between computed gradients and masks-
    for grad_layer, mask in zip(grads, mask_model.trainable_weights):
        grad_mask_mul.append(tf.math.multiply(grad_layer, mask))

    # Apply computed gradients to model's weights and biases-
    optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))

    # Compute accuracy-
    train_loss(loss)
    train_accuracy(y, y_pred)

    return None

В коде: «mask_model» - это маска, которая имеет значение 0 или 1. Использование «mask_model» предназначено для контроля того, какие параметры обучаются (так как 0 * градиентный спуск = 0).

У меня вопрос, я использую Переменная списка "grad_mask_mul" внутри "train_one_step ()" оформленной функции TensorFlow. Может ли это вызвать какие-либо проблемы, такие как:

ValueError: функция с декорированием tf.function пыталась создать переменные при первом вызове.

Или вы, ребята, видите какая-то проблема использования списочной переменной внутри функции с тензорным потоком?

Спасибо!