burn::nn::modules::attention

Function generate_padding_mask

pub fn generate_padding_mask<B>(
    pad_token: usize,
    tokens_list: Vec<Vec<usize>>,
    seq_length: impl Into<SeqLengthOption>,
    device: &<B as Backend>::Device,
) -> GeneratePaddingMask<B>where
    B: Backend,

Expand description

Generates a padding attention mask for a batch of token sequences.

§Arguments

pad_token - The token ID used for padding
tokens_list - Vector of token sequences (each sequence is a vector of token IDs)
seq_length - Sequence length option (NoMax, Max, or Fixed)
device - The device for tensor operations

§Returns

A GeneratePaddingMask containing the padded tensor and corresponding mask