Speed up training and inference: AMP, compile, fused Adam, batching

[jcwang587]

Summary

Training (cgcnn-tr / cgcnn-ft), on CUDA devices:

• bf16 torch.autocast around forward pass and loss
• model.compile() (in-place, keeps state_dict keys clean)
• fused Adam

Data movement:

• non_blocking=True on all host-to-device transfers so pin_memory actually overlaps copies with compute
• persistent_workers when DataLoader workers > 0, keeping workers and their per-process LRU caches alive across epochs

Inference:

• cgcnn_pred batches inference (new batch_size parameter, default 256) instead of hardcoded batch size 1
• cgcnn_descriptor handles batches correctly; each crystal's descriptor is now a 1D (h_fea_len,) array instead of (1, h_fea_len) - np.vstack/np.array consumers are unaffected

Cleanups: cgcnn_ft no longer builds a throwaway CIFData just to read feature dimensions; layer resets use .to(args.device); redundant .to(device) on the Boltzmann bias removed. Tutorial notebooks re-executed against the new behavior and p02's torch.load cell modernized.

Verification

• End-to-end smoke runs of cgcnn-tr, cgcnn-ft, cgcnn-pr, and the cgcnn_pred API (including the bundled pretrained model) all pass on this branch.

Part 4/4, stacked on #165 (chore/safe-checkpoint-loading).

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: ad663343-9682-43fe-9130-55377cb01e8d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}