Ускорение градиентной ленты Tensorflow 2.0

Я следовал учебному пособию TF 2.0 по сверточным VAE, расположенным здесь .

Поскольку это очень важно, градиенты вычисляются вручную, а затем применяются вручную, используя tf.GradientTape ().

for epoch in epochs:
  for x in x_train:
    with tf.GradientTape() as tape:
      loss = compute_loss(model, x)
    apply_gradients(tape.gradient(loss, model.trainable_variables))

Проблема с этим кодом в том, что он довольно медленный, занимая около 40-50 секунд на эпоху. Если я значительно увеличу размер пакета (примерно до 2048), то это займет около 8 секунд на эпоху, но производительность модели значительно снизится.

С другой стороны, если я делаю более традиционную модель (то есть, которая использует модель на основе ленивых графов вместо рвения), такую ​​как модель здесь , то это займет 8 секунд за эпоху даже с небольшим размером партии.

model.add_loss(lazy_graph_loss)
model.fit(x_train epochs=epochs)

Основываясь на этой информации, я думаю, что проблема с кодом TF2.0 заключается в ручном вычислении потерь и градиентов.

Есть ли способ ускорить код TF2.0, чтобы он был ближе к нормальному коду?