Timing to call .zero_grad()

gaozhihan / PreDiff

[NeurIPS 2023] Official implementation of "PreDiff: Precipitation Nowcasting with Latent Diffusion Models"

Apache License 2.0

80 stars 5 forks source link

Timing to call .zero_grad() #32

Open JackChan1123 opened 1 month ago

JackChan1123 commented 1 month ago

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.

gaozhihan commented 1 month ago

Thank you very much for identifying this bug. It has been fixed in #33 .