fix bugs in fused softmax

0d350c8d · hyunwoongko · mshoeybi · 116820a5 · 0d350c8d
Commit 0d350c8d authored Aug 13, 2021 by hyunwoongko Committed by mshoeybi Aug 22, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

megatron/model/fused_softmax.py megatron/model/fused_softmax.py +1 -1

No files found.
--- a/megatron/model/fused_softmax.py
+++ b/megatron/model/fused_softmax.py
@@ -186,7 +186,7 @@ class FusedScaleMaskSoftmax(nn.Module):
        return probs

    @staticmethod
-    def get_batch_per_block(b, np, sq, sk):
+    def get_batch_per_block(sq, sk, b, np):
        import scaled_masked_softmax_cuda

        return scaled_masked_softmax_cuda.get_batch_per_block(sq, sk, b, np)