What causes "Wrong shape for input_ids or attention_mask" in huggingface/candle?

Ensure the mask is 2-D [batch, seq_len] (or 3-D) before calling the model

Wrong shape for input_ids or attention_mask

Error message

Wrong shape for input_ids or attention_mask

What it means

BERT's get_extended_attention_mask only accepts attention masks of rank 2 (batch x seq_len) or rank 3, extending them for multi-head attention. Any other rank fails with this message, indicating the input_ids/attention_mask were built with an unexpected shape.

Source

Thrown at candle-transformers/src/models/bert.rs:519

        let _enter = self.span.enter();
        let embedding_output = self.embeddings.forward(input_ids, token_type_ids)?;
        let attention_mask = match attention_mask {
            Some(attention_mask) => attention_mask.clone(),
            None => input_ids.ones_like()?,
        };
        let dtype = embedding_output.dtype();
        // https://github.com/huggingface/transformers/blob/6eedfa6dd15dc1e22a55ae036f681914e5a0d9a1/src/transformers/models/bert/modeling_bert.py#L995
        let attention_mask = get_extended_attention_mask(&attention_mask, dtype)?;
        let sequence_output = self.encoder.forward(&embedding_output, &attention_mask)?;
        Ok(sequence_output)
    }
}

fn get_extended_attention_mask(attention_mask: &Tensor, dtype: DType) -> Result<Tensor> {
    let attention_mask = match attention_mask.rank() {
        3 => attention_mask.unsqueeze(1)?,
        2 => attention_mask.unsqueeze(1)?.unsqueeze(1)?,
        _ => candle::bail!("Wrong shape for input_ids or attention_mask"),
    };
    let attention_mask = attention_mask.to_dtype(dtype)?;
    // torch.finfo(dtype).min
    (attention_mask.ones_like()? - &attention_mask)?.broadcast_mul(
        &Tensor::try_from(f32::MIN)?
            .to_device(attention_mask.device())?
            .to_dtype(dtype)?,
    )
}

//https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/models/bert/modeling_bert.py#L752-L766
struct BertPredictionHeadTransform {
    dense: Linear,
    activation: HiddenActLayer,
    layer_norm: LayerNorm,
}

impl BertPredictionHeadTransform {

View on GitHub (pinned to d5fee525bf)

Solutions

Ensure the mask is 2-D [batch, seq_len] (or 3-D) before calling the model
Add a batch dimension: mask.unsqueeze(0) for a single sequence
Verify argument order so a mask isn't passed where input_ids is expected
Squeeze extra trailing dims if the mask is already 4-D from another library

Example fix

// before
let mask = Tensor::ones((128, DType::U8)...)  // rank 1
let out = model.forward(&ids, &mask, ...)?;
// after
let mask = Tensor::ones((1, 128, DType::U8)...).unsqueeze(0)?; // rank 2/3
let out = model.forward(&ids, &mask, ...)?;

Defensive patterns

Strategy: validation

Validate before calling

let rank = attention_mask.rank();
assert!(rank == 2 || rank == 3, "attention_mask rank {} unsupported; expected 2 or 3", rank);

Type guard

fn is_valid_attention_mask(mask: &Tensor) -> bool {
    matches!(mask.rank(), 2 | 3)
}

Try / catch

let mask = match attention_mask.rank() {
    2 | 3 => attention_mask,
    1 => attention_mask.unsqueeze(0)?,
    _ => return Err(anyhow!("unsupported mask rank")),
};
let out = model.forward(&input_ids, &mask, &token_type_ids, None, None)?;

Prevention

Keep masks batched: always [batch, seq_len]
Check mask rank before calling the model
Don't pass already-expanded 4-D masks from other frameworks
Verify you're not swapping input_ids and attention_mask arguments

When it happens

Trigger: Passing an attention_mask of rank 1 (single sequence, no batch dim), rank 4, or a mask already expanded to per-head shape into BertModel with attention masks.

Common situations: Preprocessing producing an unbatched 1-D mask, frameworks that already return [b, heads, seq, seq] masks, or accidentally swapping input_ids and attention_mask arguments.

Related errors

AI-assisted analysis of huggingface/candle@d5fee525bf (2026-09-02). Data as JSON: /api/errors/a1376305562a73f7. Report an issue: GitHub.