Skip to content

GitLab

Unverified Commit 6f4bc334 authored by Jack's avatar Jack Committed by GitHub
Browse files

Make grad_output contiguous in cross_entropy.py (#2402) 


Signed-off-by: default avatarJack <lityangweiguang@163.com>
parent 15dead11
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
transformer_engine/pytorch/triton/cross_entropy.py
View file @ 6f4bc334
...@@ -121,7 +121,7 @@ def cross_entropy_backward( ...@@ -121,7 +121,7 @@ def cross_entropy_backward(
element_mul_kernel[(n_rows,)]( element_mul_kernel[(n_rows,)](
_input, _input,
_input.stride(-2), _input.stride(-2),
grad_output, grad_output.contiguous(),
1 if grad_output.numel() > 1 else 0, 1 if grad_output.numel() > 1 else 0,
V, V,
BLOCK_SIZE=BLOCK_SIZE, BLOCK_SIZE=BLOCK_SIZE,
......
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