megatron.data.biencoder_dataset_utils.make_attention_mask

megatron.data.biencoder_dataset_utils.make_attention_mask(source_block, target_block)

Returns a 2-dimensional (2-D) attention mask :param source_block: 1-D array :param target_block: 1-D array