Skip to content

GitLab

Unverified Commit e1c4f51e authored by Xiaowei Ren's avatar Xiaowei Ren Committed by GitHub
Browse files

make sure dout is contiguous (#1539) 


Signed-off-by: default avatarXiaowei Ren <xren@nvidia.com>
parent 74983b36
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
transformer_engine/pytorch/attention.py
View file @ e1c4f51e
......@@ -2858,6 +2858,7 @@ class AttnFuncWithCPAndKVP2P(torch.autograd.Function):
# [b, np, sq] -> [b, np, sq, 1] or
# [t, np] -> [t, np, 1]
softmax_lse.unsqueeze_(-1)
dout = dout.contiguous()
dq = None
dout_dtype = dout.dtype
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment