What causes "Capture tensors do not support autodiff" in tracel-ai/burn?

Do not call backward() on capture tensors: run capture only in inference/export paths

Capture tensors do not support autodiff

Error message

Capture tensors do not support autodiff

What it means

burn-dispatch's AutodiffBackend::backward() calls .backward() on the tensor's inner backend kind. The Capture backend (operation capture for export/compilation) has no autodiff implementation, so any backward() on a Capture-kind tensor is an immediate panic. This is an intentional unsupported-operation guard, not a recoverable failure.

Source

Thrown at crates/burn-dispatch/src/backend.rs:345

    fn backward(tensor: DispatchTensor) -> Self::Gradients {
        let DispatchTensor { kind, .. } = tensor;

        match kind {
            DispatchTensorKind::Autodiff(tensor) => match *tensor {
                #[cfg(cube_backend)]
                DispatchTensorKind::Cube(tensor) => tensor.autodiff().backward(),
                #[cfg(any(feature = "flex", default_backend))]
                DispatchTensorKind::Flex(tensor) => tensor.autodiff().backward(),
                #[cfg(feature = "ndarray")]
                DispatchTensorKind::NdArray(tensor) => tensor.autodiff().backward(),
                #[cfg(feature = "tch")]
                DispatchTensorKind::LibTorch(tensor) => tensor.autodiff().backward(),
                #[cfg(feature = "remote")]
                DispatchTensorKind::Remote(tensor) => tensor.autodiff().backward(),
                #[cfg(feature = "capture")]
                DispatchTensorKind::Capture(_) => {
                    panic!("Capture tensors do not support autodiff")
                }
                DispatchTensorKind::Autodiff(_) => {
                    panic!("Autodiff should not wrap an autodiff tensor.")
                }
            },
            _ => panic!("Requires autodiff tensor."),
        }
    }

    fn grad(tensor: &DispatchTensor, grads: &Self::Gradients) -> Option<DispatchTensor> {
        let DispatchTensor { kind, .. } = tensor;
        let grad: Option<DispatchTensorKind> = match &kind {
            DispatchTensorKind::Autodiff(inner_kind) => match &**inner_kind {
                #[cfg(cube_backend)]
                DispatchTensorKind::Cube(tensor) => tensor
                    .as_autodiff()
                    .grad(grads)
                    .map(|t| DispatchTensorKind::Cube(crate::BackendTensor::Float(t))),

View on GitHub (pinned to d16f7ba2ed)

Solutions

Do not call backward() on capture tensors: run capture only in inference/export paths
Switch the tensor/model to a real backend (NdArray, LibTorch, Cube, Flex, Remote) before training
Guard training code with a backend check so capture-mode tensors never reach backward()

Example fix

// before
let grads = loss.backward(); // panics: loss is a Capture tensor
// after
let grads = match loss_kind {
    DispatchTensorKind::Capture(_) => return Err(Error::CaptureNoAutodiff),
    _ => loss.backward(),
};

Defensive patterns

Strategy: validation

Validate before calling

fn ensure_not_capture(t: &DispatchTensor) -> Result<(), String> {
    if matches!(t.kind, DispatchTensorKind::Capture(_)) {
        Err("capture tensors do not support autodiff; use a real backend for training".into())
    } else { Ok(()) }
}

Type guard

fn is_capture(t: &DispatchTensor) -> bool {
    matches!(t.kind, DispatchTensorKind::Capture(_))
}

Try / catch

// Rust panics are not catchable with try/catch; use catch_unwind if unavoidable
let result = std::panic::catch_unwind(AssertUnwindSafe(|| Dispatch::backward(loss)));
match result {
    Ok(grads) => use_grads(grads),
    Err(_) => eprintln!("capture tensors do not support autodiff"),
}

Prevention

Never mix the capture feature into training/autodiff code paths
Assert tensor kind before any AutodiffBackend call
Keep capture/export passes separate from training passes

When it happens

Trigger: Calling Dispatch::backward(tensor) (directly or via a training loop / loss.backward()) where the tensor was created under the 'capture' feature and is a DispatchTensorKind::Capture, e.g. building a model for burn-export/capture then attempting gradient computation.

Common situations: Using the capture backend (tensor capture for ONNX/kernel export) while accidentally running in a training/autodiff context; enabling both 'capture' and 'autodiff' dispatch features and routing a loss tensor through capture.

Related errors

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/98e3da6c389b4a75. Report an issue: GitHub.