Added Diffentiable per_sample_weights Check to EmbeddingBag.cpp #142338
Conversation
Author: Emmett Bicker Added a check in aten/src/ATen/native/EmbeddingBag.cpp that checks if per_sample_weights needs a gradient in order to determine if forward_only or the full embedding bag implementation should run. Also, added two tests in test_embedding.py that check if the command now works.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/142338
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 7a54d87 with merge base 524395e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
test/nn/test_embedding.py
Outdated
| x = torch.arange(1, 5, device=device).expand(3, -1) | ||
| w = torch.rand(3, 4, device=device, requires_grad=per_sample_weights_use_grad) | ||
| bag(x, per_sample_weights=F.softmax(w, dim=-1)) | ||
|
|
There was a problem hiding this comment.
Can you put this test under TestEmbeddingNNDeviceType? This lets us parametrize on device as well /avoid duplication between cpu/cuda.
There was a problem hiding this comment.
Sure! Ill update the PR shortly
soulitzer
added
release notes: autograd
janeyx99
added
the
triaged
|
@soulitzer Does this look better? |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 3 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
@soulitzer Hi! Sorry I think I removed some whitespace at the same time you @ merged and it might have broken it? Which would make a lot of sense. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
Hi @soulitzer ! I'm so sorry I'm still getting used to contributing and forgot to run the linter on the other file I edited, I just removed the two other problematic pieces of whitespace in another commit. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
Woohoo! |
5 participants
Added a check in aten/src/ATen/native/EmbeddingBag.cpp that checks if per_sample_weights needs a gradient in order to determine if at::_embedding_bag_forward_only or at::_embedding_bag should run.
Also, added two tests in test_embedding.py that check if the command now works.
Fixes #136457