Timing to call .zero_grad()

I noticed that you call .zero_grad() outside the IF condition in your VAE training script (line 454 & 467). Shouldn't we call it after .step() in IF condition? Otherwise the accumulated gradient in computation graphs will be erased.