huggingface/candle · error · candle::Error

the target tensor should have two dimensions ({dims:?})

Error message

the target tensor should have two dimensions ({dims:?})

What it means

nll requires the input (already log-softmaxed logits) to be exactly 2-D of shape [batch, num_classes]. This bail fires for any other rank — e.g. 3-D sequence logits of shape [B, T, C] — because the single gather is written for the 2-D classification case only.

Source

Thrown at candle-nn/src/loss.rs:25

/// Arguments
///
/// * [inp]: The input tensor of dimensions `N, C` where `N` is the batch size and `C` the number
///   of categories. This is expected to contain log probabilities.
/// * [target]: The ground truth labels as a tensor of u32 of dimension `N`.
///
/// The resulting tensor is a scalar containing the average value over the batch.
pub fn nll(inp: &Tensor, target: &Tensor) -> Result<Tensor> {
    let b_sz = match target.dims() {
        &[b_sz] => b_sz,
        dims => candle::bail!("the target tensor should have a single dimension ({dims:?})"),
    };
    match inp.dims() {
        &[inp_b_sz, _] => {
            if inp_b_sz != b_sz {
                candle::bail!("batch size mismatch between inp ({inp_b_sz}) and target ({b_sz})")
            }
        }
        dims => candle::bail!("the target tensor should have two dimensions ({dims:?})"),
    }
    inp.gather(&target.unsqueeze(1)?, 1)?
        .sum_all()?
        .affine(-1f64 / b_sz as f64, 0.)
}

/// The cross-entropy loss.
///
/// Arguments
///
/// * [inp]: The input tensor of dimensions `N, C` where `N` is the batch size and `C` the number
///   of categories. This is expected to raw logits.
/// * [target]: The ground truth labels as a tensor of u32 of dimension `N`.
///
/// The resulting tensor is a scalar containing the average value over the batch.
pub fn cross_entropy(inp: &Tensor, target: &Tensor) -> Result<Tensor> {
    if inp.rank() != 2 {
        candle::bail!("cross_entropy expects an input tensor of rank 2")

View on GitHub (pinned to d5fee525bf)

Solutions

  1. For sequence logits, reshape [B, T, C] to [B*T, C] and flatten targets to [B*T] before calling cross_entropy, then reshape the scalar/mean as needed.
  2. Add a batch dimension for unbatched input: logits.unsqueeze(0)?.
  3. For per-token losses, reshape to 2-D, compute cross_entropy, and apply your own masking/averaging over valid tokens.
  4. Use mse_loss/other losses only for regression; for structured outputs write a custom NLL with gather on the appropriate axis.

Example fix

// before: logits [B, T, C]
let loss = cross_entropy(&logits, &targets)?; // error
// after
let (b, t, _c) = logits.dims3()?;
let loss = cross_entropy(&logits.reshape((b * t, ()))?, &targets.reshape((b * t, ()))?)?;
Defensive patterns

Strategy: validation

Validate before calling

if inp.rank() != 2 {
    return Err(anyhow!("nll input must be [B, C], got dims {:?}", inp.dims()));
}

Type guard

fn is_logit_matrix(x: &Tensor) -> bool { x.rank() == 2 }

Try / catch

let (logits, targets) = if logits.rank() == 3 {
    let (b, t, _c) = logits.dims3()?;
    (logits.reshape((b * t, ()))?, targets.reshape((b * t, ()))?)
} else { (logits, targets) };
let loss = cross_entropy(&logits, &targets)?;

Prevention

When it happens

Trigger: Calling nll or cross_entropy with inp of rank 1, 3, or higher — commonly per-token logits [B, T, C] from a language model, or [C] unbatched logits.

Common situations: Sequence/token-level classification fed straight into cross_entropy instead of flattening time; unbatched inference tensor; porting PyTorch code where F.cross_entropy supports [N, C, d1..] and assuming candle does too.

Related errors


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