pool

hypembed::model::pool

Function pool 

Source 
pub fn pool(
    hidden: &Tensor,
    attention_mask: &[Vec<u32>],
    strategy: PoolingStrategy,
) -> Result<Tensor>
Expand description

Apply pooling to hidden states.

§Arguments

  • hidden: [batch, seq, hidden_size]
  • attention_mask: [batch, seq] with 1 = real and 0 = padding
  • strategy: Which pooling to use

§Returns

Pooled tensor [batch, hidden_size]