Skip to content

fix(video_jepa): apply AMP from config in training loop - #19

Open
Delata22 wants to merge 1 commit into
facebookresearch:mainfrom
Delata22:fix/video-jepa-amp
Open

fix(video_jepa): apply AMP from config in training loop#19
Delata22 wants to merge 1 commit into
facebookresearch:mainfrom
Delata22:fix/video-jepa-amp

Conversation

@Delata22

Copy link
Copy Markdown

The video_jepa training config declares use_amp=true with dtype=float16, but main.py never read cfg.training, so training ran in fp32 regardless of the config. This is inconsistent with image_jepa and ac_video_jepa, which both apply AMP correctly.

This applies the same autocast/GradScaler pattern used in those examples. Since detection_head uses nn.BCELoss (unsafe under autocast), det_loss is computed outside the autocast block.

The video_jepa training config declares use_amp=true with
dtype=float16, but main.py never read cfg.training, so training ran
in fp32 regardless of the config. This is inconsistent with
image_jepa and ac_video_jepa, which both apply AMP correctly.

This applies the same autocast/GradScaler pattern used in those
examples. Since detection_head uses nn.BCELoss (unsafe under
autocast), det_loss is computed outside the autocast block.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label May 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant