Unfortunately, using save_for_backward does not change the figures anyhow, so the overhead is still there. Here is the updated code:
https://colab.research.google.com/drive/1iGZJrNGu507V2QmpSEO6tQpCK61sC8jD 5