code: def attention_mask_func(attention_scores, attention_mask):
      attention_scores.masked_fill_(attention_mask, -10000.0)
      return attention_scores
  
file_path: ColossalAI/examples/tutorial/sequence_parallel/model/layers/bert_layer.py
imports: import torch
  import torch.nn
  from colossalai.kernel.cuda_native import LayerNorm
  from colossalai.kernel.jit import bias_dropout_add_fused_inference, bias_dropout_add_fused_train
  from colossalai.legacy.nn.layer.parallel_sequence import TransformerSelfAttentionRing
  from .dropout import get_bias_dropout_add
  from .mlp import TransformerMLP
label_desc: Use image node_id 1 for calling a global function with example usage: attention_mask_func(attention_scores, attention_mask) and returns: attention_scores