Fix: KronosPredictor doesn't call .eval(), breaking inference on MPS backend by Lu-Derik · Pull Request #407 · shiyu-coder/Kronos · GitHub
Skip to content

Fix: KronosPredictor doesn't call .eval(), breaking inference on MPS backend - #407

Open
Lu-Derik wants to merge 1 commit into
shiyu-coder:masterfrom
Lu-Derik:fix/mps-inference-eval-mode
Open

Lu-Derik wants to merge 1 commit into
shiyu-coder:masterfrom
Lu-Derik:fix/mps-inference-eval-mode

Conversation

@Lu-Derik

Copy link
Copy Markdown

Summary

  • KronosPredictor.__init__ moves the tokenizer/model to device but never calls .eval(), leaving them in PyTorch's default training=True mode during inference.
  • This keeps dropout active inside F.scaled_dot_product_attention, which raises NotImplementedError: scaled_dot_product_attention for MPS does not support dropout on Apple Silicon (MPS backend). CUDA silently tolerates dropout in SDPA, so this went unnoticed there.

Fix

Call .eval() on both the tokenizer and model right after moving them to device.

Test plan

  • Ran examples/prediction_example.py on macOS (MPS backend) — previously crashed with the error above, now completes successfully end-to-end.

🤖 Generated with Claude Code

KronosPredictor moved the tokenizer/model to device but never called
.eval(), so PyTorch's default training=True mode stayed active during
inference. This kept dropout enabled inside
F.scaled_dot_product_attention, which raises
"NotImplementedError: scaled_dot_product_attention for MPS does not
support dropout" on Apple Silicon (MPS backend). CUDA silently
tolerates dropout in SDPA, which is why this went unnoticed there.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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.

1 participant