Upload GPTJXForCausalLM

pretrained_model.py

@@ -58,6 +58,7 @@ class CausalSelfAttention(nn.Module):
         if self.flash:
             if attn_mask is not None:
             # efficient attention using Flash Attention CUDA kernels
+                attn_mask = attn_mask.to(torch.bool)
                 y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0)
             else:
                 y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)