Fix checkpoint saving using wrong interval config

[Mr-Neutr0n]

Summary

The checkpoint saving condition in train.py (line 242) has a bug where it checks config.checkpoint_every to determine whether checkpointing is enabled, but then incorrectly uses config.eval_every for the modulo operation:

# Before (bug):
if ((config.checkpoint_every and step % config.eval_every == 0) or
    step == total_steps):

# After (fix):
if ((config.checkpoint_every and step % config.checkpoint_every == 0) or
    step == total_steps):

This causes checkpoints to be saved at evaluation intervals instead of the configured checkpoint intervals, making the checkpoint_every config value effectively ignored for controlling save frequency.

Test plan

• Verified the fix by inspecting the changed line
• The fix is a single-token change from config.eval_every to config.checkpoint_every in the modulo check, matching the guard condition on the same line

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[Mr-Neutr0n]

The only CI failure here is cla/google — the Google CLA check. The code change itself is correct and the diff is minimal (1 line in vit_jax/train.py).

This is a human action: the CLA needs to be signed at https://cla.developers.google.com/ before Google can accept the PR. The check should re-run automatically once the CLA is on file for the Mr-Neutr0n GitHub account. No code changes needed from my side.

Leaving the PR open — the maintainers will pick it up once the CLA clears.