Checkpoint files missing when `use_infoNCE` is set to `True`

[Xinming-Dai]

Describe the bug
When use_infoNCE: True is specified in the vit.yaml configuration file, the training run completes but no .ckpt files are produced in the output directory. This differs from the expected behavior observed when use_infoNCE: False, where checkpoint files are correctly saved.

To Reproduce

• Open vit.yaml.
• Change use_infoNCE: False to use_infoNCE: True.
• Run beast train.
• Check the output directory.

Expected behavior
The output directory should contain a checkpoints/ folder with the corresponding .ckpt files, consistent with runs where use_infoNCE is set to False.

[themattinthehatt]

Huh that is very strange...I was able to train a model with use_infoNCE a couple months ago without issue.

I'm quite swamped with ICLR reviews right now but can dive into this more on Wednesday. If you happen to track down the issue in the meantime please keep me updated!

(And thanks for the bug finds, always good to have multiple pairs of eyes on the code 🙏)