Hi @hduc-le, there is a new context manager allow_mutation_on_saved_mode.py · GitHub is designed to get rid of these “variables needed for gradient computation has been modified” errors. I have tried it on your colab and it no longer runs into the error. Let me know how it works for you.
The usage is
with allow_mutation_on_saved_tensors():
# all the code in your training loop
...