[PYTHON] Bugfix in trans_c for dsd in blocksparse matmul (#131)

2021-07-12 23:57:44 -07:00
parent 8cea583109
commit 3169e4355c
1 changed files with 2 additions and 2 deletions
--- a/python/triton/ops/blocksparse/matmul.py
+++ b/python/triton/ops/blocksparse/matmul.py
@@ -512,8 +512,8 @@ class _matmul(torch.autograd.Function):
            b.stride(2 if trans_b else 3),
            c.stride(0),
            c.stride(1),
-            c.stride(2),
-            c.stride(3),
+            c.stride(3 if trans_c else 2),
+            c.stride(2 if trans_c else 3),
            BS3,
            AS1,
            0,