Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix bug: noise model not training on GPU #360

Merged
merged 2 commits into from
Jan 20, 2025
Merged

Conversation

veegalinova
Copy link
Collaborator

Description

Fixes a bug where the noise model was unintentionally moved to the CPU before training, and thus did not train on GPU.

  • What: When _set_model_mode was called, the model was moved to the CPU by default, which caused training to happen on the CPU instead of the GPU.

Changes Made

  • Modified: Removed device handling from _set_model_mode, moved final move to CPU to the end of the fit function.

Please ensure your PR meets the following requirements:

  • Code builds and passes tests locally, including doctests
  • New tests have been added (for bug fixes/features)
  • Pre-commit passes
  • PR to the documentation exists (for bug fixes / features)

@jdeschamps jdeschamps merged commit cf3a094 into main Jan 20, 2025
18 checks passed
@jdeschamps jdeschamps deleted the vg/fix/noise-model-device branch January 20, 2025 14:58
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.

3 participants