Skip to content

Commit

Permalink
Fix cuda/cpu check on NoneType (pytorch#88854)
Browse files Browse the repository at this point in the history
Summary: Fix cuda/cpu check on NoneType

Test Plan: sabdcastle/ github CI/CD

Differential Revision: D41203955

Pull Request resolved: pytorch#88854
Approved by: https://github.com/drisspg, https://github.com/ngimel
  • Loading branch information
mikekgfb authored and weiwangmeta committed Dec 6, 2022
1 parent a81f9b3 commit ba9d89b
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion torch/nn/modules/activation.py
Expand Up @@ -1116,7 +1116,7 @@ def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: O
# generator expressions.
if torch.overrides.has_torch_function(tensor_args):
why_not_fast_path = "some Tensor argument has_torch_function"
elif not all([(x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args]):
elif not all([(x is None or x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args]):
why_not_fast_path = "some Tensor argument is neither CUDA nor CPU"
elif torch.is_grad_enabled() and any([x.requires_grad for x in tensor_args]):
why_not_fast_path = ("grad is enabled and at least one of query or the "
Expand Down

0 comments on commit ba9d89b

Please sign in to comment.