Skip to content

Fix predict_on_seqs OOM via batched Trainer.predict - #195

Open
buyan-kh wants to merge 1 commit into
Genentech:mainfrom
buyan-kh:fix/predict-on-seqs-oom
Open

buyan-kh wants to merge 1 commit into
Genentech:mainfrom
buyan-kh:fix/predict-on-seqs-oom

Conversation

@buyan-kh

@buyan-kh buyan-kh commented Sep 8, 2026

Copy link
Copy Markdown

Summary

  • Fixes Weird OOM in 1_inference.ipynb #154: predict_on_seqs eagerly moved the full one-hot batch onto the device and ran a single forward pass, which OOMs on large models (e.g. Borzoi in 1_inference.ipynb on 16GB VRAM).
  • Aligns predict_on_seqs with predict_on_dataset / ISM_predict by using make_predict_loader + Lightning Trainer.predict so inference is batched.
  • Adds optional devices, num_workers, batch_size, and precision kwargs (defaults: batch_size=1 for memory safety on large models).
  • Keeps the device= keyword as a backward-compatible alias so existing tutorials (predict_on_seqs(..., device=0)) continue to work.

Test plan

  • pytest tests/test_lightning.py::test_lightning_model_predict_on_seqs (CPU; covers single/multi seqs, devices=, device= alias, and batch_size)
  • pytest tests/test_lightning.py::test_lightning_model_predict_on_dataset still passes
  • Optional: re-run Borzoi predict_on_seqs from docs/tutorials/1_inference.ipynb on a 16GB GPU

Align predict_on_seqs with predict_on_dataset / ISM_predict by routing
through make_predict_loader and Lightning Trainer.predict instead of
eagerly moving the full one-hot batch onto the device. Keeps the device=
keyword for backward compatibility with existing tutorials.

Fixes Genentech#154
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Weird OOM in 1_inference.ipynb

1 participant