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

Remove CUDA synchronizations by slicing input tensor with int instead of CUDA tensors in nn.LinearEmbeddingEncoder #432

Merged
merged 4 commits into from
Aug 12, 2024

Conversation

akihironitta
Copy link
Member

start_idx and end_idx used at feat.values[:, start_idx:end_idx] are on device, which leads to cuda synchronizations.

@akihironitta akihironitta changed the title Remove CUDA synchronizations by slicing input tensor with int instead of CUDA tensors Remove CUDA synchronizations by slicing input tensor with int instead of CUDA tensors in nn.LinearEmbeddingEncoder Aug 10, 2024
@akihironitta akihironitta requested a review from rusty1s August 10, 2024 15:13
Copy link
Contributor

@yiweny yiweny left a comment

Choose a reason for hiding this comment

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

Good catch. Thank you!

torch_frame/nn/encoder/stype_encoder.py Show resolved Hide resolved
@akihironitta akihironitta merged commit 1f4c4b8 into master Aug 12, 2024
13 checks passed
@akihironitta akihironitta deleted the aki-rm-cuda-sync branch August 12, 2024 09:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants