Wrong Implementation of perturb_input

Hi
I think there is an error in the notebook in the implementation of perturb_input, the weight for the noise should be square-rooted.

Current:
def perturb_input(x, t, noise):
return ab_t.sqrt()[t, None, None, None] * x + (1 - ab_t[t, None, None, None]) * noise

Expected:
def perturb_input(x, t, noise):
return ab_t.sqrt()[t, None, None, None] * x + (1 - ab_t[t, None, None, None]).sqrt() * noise

This will make it aligned with the original paper https://arxiv.org/pdf/2006.11239.pdf, and also with the implementation in minDiffusion Repo https://github.com/cloneofsimo/minDiffusion/blob/master/mindiffusion/ddpm.py#L78

Please correct me if I am wrong, or let know the course instructors to make a fix.

Thank :slight_smile: