Skip to content

torch.load always to cpu first#308

Merged
johannaSommer merged 1 commit intomainfrom
fix/torch-save-always-load-to-cpu
Aug 20, 2025
Merged

torch.load always to cpu first#308
johannaSommer merged 1 commit intomainfrom
fix/torch-save-always-load-to-cpu

Conversation

@simlang
Copy link
Copy Markdown
Member

@simlang simlang commented Aug 19, 2025

Description

This PR does a little fix
When using torch.load, always load to cpu first and then move to target device.
Problem when loading to the device directly was, that every tensor within a pipeline e.g. was moved to cuda, also some which are not supposed to be - not every pipeline can handle this unexpected device (e.g. SanaPipeline)
By loading to cpu first and then moving to target device, we make sure that only the right parts are moved to the target_device

Related Issue

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

How Has This Been Tested?

Every test which is using torch.load has been re-run and is successful

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

Additional Notes

Copy link
Copy Markdown
Collaborator

@gsprochette gsprochette left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, the fix itself is ready to merge but could you add a test with Sana so we can see the before-after, and so the reason for this "inefficient" loading becomes clear to someone who would like to optimize it when they run the tests?

@simlang
Copy link
Copy Markdown
Member Author

simlang commented Aug 20, 2025

there are already tests with Sana, which are currently failing - this PR fixes these tests

@simlang simlang requested a review from gsprochette August 20, 2025 11:13
Copy link
Copy Markdown
Collaborator

@gsprochette gsprochette left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

True. LGTM then :)

Copy link
Copy Markdown
Member

@johannaSommer johannaSommer left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🫡 merci

@johannaSommer johannaSommer merged commit 4582abc into main Aug 20, 2025
7 checks passed
@johannaSommer johannaSommer deleted the fix/torch-save-always-load-to-cpu branch August 20, 2025 14:33
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