facebookresearch / ijepa

Official codebase for I-JEPA, the Image-based Joint-Embedding Predictive Architecture. First outlined in the CVPR paper, "Self-supervised learning from images with a joint-embedding predictive architecture."
Other
2.83k stars 358 forks source link

Why there is no **unscale_** when you use amp? #64

Open jsrdcht opened 5 months ago

jsrdcht commented 5 months ago

Your code is

if use_bfloat16:
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

I think it should be

if use_bfloat16:
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    scaler.step(optimizer)
    scaler.update()

Am I right?