Skip to content

fix: set current CUDA device in _inplace_pin_memory function#77

Merged
blahgeek merged 1 commit intoMoonshotAI:mainfrom
SongXiaoXi:main
Jan 9, 2026
Merged

fix: set current CUDA device in _inplace_pin_memory function#77
blahgeek merged 1 commit intoMoonshotAI:mainfrom
SongXiaoXi:main

Conversation

@SongXiaoXi
Copy link
Collaborator

This pull request introduces a minor import adjustment and a fix to device handling when pinning memory for CUDA tensors. The main changes ensure that memory pinning happens on the correct CUDA device and resolve an import path issue.

CUDA device handling improvements:

  • In checkpoint_engine/pin_memory.py, the current CUDA device is now explicitly retrieved and set (device_index = torch.cuda.current_device() and torch.cuda.set_device(device_index)) before pinning memory, ensuring that pinning occurs on the correct device. [1] [2]

Import path correction:

  • In examples/update.py, the import for request_inference_to_update is fixed by importing it directly from checkpoint_engine instead of checkpoint_engine.ps, resolving a potential import error.

@blahgeek blahgeek merged commit 4a73109 into MoonshotAI:main Jan 9, 2026
2 checks passed
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.

2 participants