huggingface/candle · error · candle::Error

cross_entropy expects an input tensor of rank 2

Error message

cross_entropy expects an input tensor of rank 2

What it means

cross_entropy is a thin wrapper that applies log_softmax then nll, and it only supports rank-2 inputs ([batch, classes]). This guard fires before any computation when inp.rank() != 2, giving an earlier and clearer message than the downstream nll shape errors.

Source

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

        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")
    }
    let inp = crate::ops::log_softmax(inp, 1)?;
    nll(&inp, target)
}

/// The mean squared error loss.
pub fn mse(inp: &Tensor, target: &Tensor) -> Result<Tensor> {
    (inp - target)?.sqr()?.mean_all()
}

/// The binary cross-entropy with logit 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, C` where `N` is the batch size and `C` the number
///   of categories.

View on GitHub (pinned to d5fee525bf)

Solutions

  1. Reshape higher-rank logits to [N, C]: e.g. [B,C,H,W] -> permute to [B,H,W,C] then reshape ((B*H*W), C), flattening targets likewise.
  2. For [B,T,C], reshape to [B*T, C] and flatten targets to [B*T].
  3. Unsqueeze a batch dim for rank-1 input before calling.
  4. If you regularly need higher-rank cross entropy, write a small wrapper that reshapes, calls cross_entropy, and reshapes back.

Example fix

// before: dense logits [B, C, H, W]
let loss = cross_entropy(&logits, &mask)?; // error
// after
let (b, c, h, w) = logits.dims4()?;
let logits = logits.permute((0, 2, 3, 1))?.reshape((b * h * w, c))?;
let mask = mask.reshape((b * h * w, ()))?;
let loss = cross_entropy(&logits, &mask)?;
Defensive patterns

Strategy: validation

Validate before calling

if logits.rank() != 2 {
    anyhow::bail!("cross_entropy requires rank 2 logits, got {:?}", logits.dims());
}

Type guard

fn is_cross_entropy_ready(logits: &Tensor, targets: &Tensor) -> bool {
    logits.rank() == 2 && targets.rank() == 1 && logits.dim(0).map(|d| d == targets.elem_count()).unwrap_or(false)
}

Try / catch

fn ce_any_rank(logits: &Tensor, targets: &Tensor) -> candle::Result<Tensor> {
    let (n, c) = match logits.rank() {
        2 => logits.dims2()?,
        _ => {
            let c = *logits.dims().last().unwrap();
            let n = logits.elem_count() / c;
            (n, c)
        }
    };
    candle_nn::loss::cross_entropy(
        &logits.reshape((n, c))?,
        &targets.reshape((targets.elem_count(), ()))?,
    )
}

Prevention

When it happens

Trigger: Calling candle_nn::loss::cross_entropy with a tensor of rank 1 (unbatched [C]), rank 3 ([B,T,C] sequence logits), or rank 4 ([B,C,H,W] dense prediction logits).

Common situations: Segmentation/dense-prediction heads outputting [B,C,H,W]; transformer token logits [B,T,C]; unbatched single sample; porting PyTorch F.cross_entropy which accepts higher-rank inputs channel-first.

Related errors

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