DDSP VST Model training & autoencoder training both instantly start to slow down significantly

[olaviinha]

Describe the bug
Training starts to slow down instantly after you start it. After a few thousand steps the pace is already considerably over 1 second per step. Around 8k steps it is already taking several seconds for each step. "Training typically takes ~2-3 hours with free Colab" mentioned in the VST Training notebook is not true at all with this bug.

System Info

• Colab Pro on T4

To Reproduce

1. Run Train VST notebook
2. Observe the training pace after a few thousand steps (e.g. around 5000 steps) -> is extremely slow.
3. Interrupt training.
4. Restart/continue training.
5. Observe the training pace after continuing training from the last saved checkpoint -> is fast again. You can observe the difference clearly with naked eye, no need to add measurements.

Additional context
Stopping and restarting !ddsp_run programmatically e.g. every 1000 steps seems to work around the issue. Otherwise it will take literally days or weeks to train a model.