Fixed an enum value comparision error (bigscience-workshop#132)

adammoody · Jun 8, 2023 · 7491937 · 7491937
1 parent d5c822e
commit 7491937
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/megatron/model/fused_softmax.py b/megatron/model/fused_softmax.py
@@ -137,14 +137,14 @@ def forward(self, input, mask):
             custom_kernel_constraint and self.scaled_masked_softmax_fusion:
             scale = self.scale if self.scale is not None else 1.0
 
-            if self.attn_mask_type == AttnMaskType.causal:
+            if self.attn_mask_type.value == AttnMaskType.causal.value:
                 assert query_seq_len == key_seq_len, \
                     "causal mask is only for self attention"
                 input = input.view(-1, query_seq_len, key_seq_len)
                 probs = ScaledUpperTriangMaskedSoftmax.apply(input, scale)
                 probs = probs.view(*data_size)
             else:
-                assert self.attn_mask_type == AttnMaskType.padding
+                assert self.attn_mask_type.value == AttnMaskType.padding.value
                 probs = ScaledMaskedSoftmax.apply(input, mask, scale)
         else:
             if self.input_in_float16 and self.softmax_in_fp32: