Resolved issue saving & loading the predictor (Regressor/Classifier) model on different devices (cuda / cpu) #486
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
… cuda device and loaded on a cpu device. Fix consists of always loading the device to cpu before saving. Reason: Once a "cuda tensor" is saved and later loaded on a cpu only machine the joblib.load() throws an error. One could monkey patch the torch.load(...,map_location="cpu") and restore it later. But moving all tensors to cpu once during saving is cleaner and has minimal one-time performance costs (+-10 ms) (and is a robust fix, without potential side-effects). also extended the "test_save_load_happy_path" test, with a `loading_device` and `saving_device`.
… on a cuda device and lader loading on a cpu-device. Also added test for different cuda-cpu saving-loading combinations.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request effectively resolves the issue of saving and loading models across different devices (CPU/CUDA) by ensuring all tensor-like objects are moved to the CPU before serialization. The renaming of save_state_expect_model_weights to save_state_except_model_weights is a good clarification. The addition of comprehensive tests for various device combinations is a great improvement. I have a few suggestions to enhance the robustness of the implementation and the tests.
…loading.py checked for tensors and nn.modules
noahho
reviewed
Sep 13, 2025
noahho
approved these changes
Sep 13, 2025
Collaborator
noahho
left a comment
noahho
left a comment
There was a problem hiding this comment.
LGTM, thanks for the work. Please only review the comment made.
Co-authored-by: noahho <noah@priorlabs.ai>
oscarkey
pushed a commit
that referenced
this pull request
Nov 12, 2025
…ressor/Classifier) model on different devices (cuda / cpu) (#156) * Record copied public PR 486 * Resolved issue saving & loading the predictor (Regressor/Classifier) model on different devices (cuda / cpu) (#486) Co-authored-by: noahho <noah@priorlabs.ai> --------- Co-authored-by: mirror-bot <mirror-bot@users.noreply.github.com> Co-authored-by: MagnusBuehler <111045718+MagnusBuehler@users.noreply.github.com> Co-authored-by: noahho <noah@priorlabs.ai>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Motivation and Context
Issue #416 reported a crash when a model saved on a CUDA device was later loaded on a CPU-only machine, due to tensors being serialized as CUDA tensors. This has been fixed by ensuring all tensor-like elements of the predictor are saved as CPU tensors, allowing safe loading on any device. The desired device can be specified (or automatically inferred) at loading time. This adds a negligible one-time overhead (~±10 ms, tested over 10 runs). Corresponding tests now cover all combinations of fitting, saving, and loading devices.
Further I added uv.lock to the .gitignore file and renamed a function in "src/tabpfn/inference.py" from
save_state_expect_model_weightstosave_state_except_model_weights, as it makes more sense with what the function is doing.Public API Changes
How Has This Been Tested?
On a cuda-environment & locally. Also added corresponding pytests.
Checklist
CHANGELOG.md(if relevant for users).