Skip to content

Conversation

@willdumm
Copy link
Contributor

Fixes #119 keeping DXSM datasets stored on CPU, except when batches are used in training.
to method will only change self.device, but will not move data between devices. The data of a dataset can be moved between devices with the move_data_to_device method, which must be implemented by all subclasses of DXSMDataset. This method will only be called in __post_init__ to ensure that all data is on the cpu. Any methods which should be executed on self.device should now be responsible for moving the required data to that device.

@willdumm willdumm marked this pull request as ready for review February 19, 2025 21:59
@willdumm willdumm requested a review from matsen February 19, 2025 21:59
@willdumm willdumm force-pushed the 119-datasets-on-cpu branch from 4ce3aba to 42ee21c Compare February 20, 2025 17:00
Copy link
Contributor

@matsen matsen left a comment

Choose a reason for hiding this comment

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

Nice!

@willdumm willdumm merged commit 2194ae4 into main Feb 22, 2025
2 checks passed
@willdumm willdumm deleted the 119-datasets-on-cpu branch February 22, 2025 21:57
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.

Keep datasets on CPU

3 participants