WGAN-GP slow critic training time
This part calculates a second derivative. it should be slow and time consuming. 

# Perform weight update
with tf.GradientTape() as outer_tape:
    # Calculate gradient penalty loss
    with tf.GradientTape() as inner_tape:


Try 

# Perform weight update
with tf.GradientTape() as outer_tape, tf.GradientTape() as inner_tape: