Device::gradient_checkpointing requires autodiff; call Devic

Error message

Device::gradient_checkpointing requires autodiff; call Device::autodiff first

What it means

Burn's Device wrapper can enable gradient checkpointing only on an autodiff device (DispatchDevice::Autodiff). Calling Device::gradient_checkpointing on a plain backend device panics with this message, because checkpointing is a training-time autodiff feature with no meaning without a backward pass. The intended workflow is device.autodiff() first, then gradient_checkpointing().

Source

Thrown at crates/burn-tensor/src/device.rs:606

    ///
    /// ```rust,ignore
    /// let device = Device::default().autodiff().gradient_checkpointing();
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if autodiff is not enabled on this device.
    #[cfg(feature = "autodiff")]
    #[must_use]
    pub fn gradient_checkpointing(self) -> Self {
        match self.into_dispatch() {
            DispatchDevice::Autodiff(device) => {
                Self::new(DispatchDevice::autodiff_with_gradient_checkpointing(
                    device.inner(),
                    GradientCheckpointingStrategy::Balanced,
                ))
            }
            _ => panic!(
                "Device::gradient_checkpointing requires autodiff; call Device::autodiff first"
            ),
        }
    }

    /// Returns this device without its autodiff association.
    ///
    /// If autodiff is not enabled, the device is returned unchanged. This operation is idempotent.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let device = Device::default().autodiff();
    /// let inference_device = device.without_autodiff();
    ///
    /// assert!(!inference_device.is_autodiff());
    /// ```
    #[must_use]

View on GitHub (pinned to d16f7ba2ed)

Solutions

Wrap the backend in autodiff first: let dev = device.autodiff(); then call dev.gradient_checkpointing()
Verify the device you call gradient_checkpointing on is the Autodiff-wrapped instance, not the raw backend device
If you only need inference, remove the gradient_checkpointing call entirely — it is meaningless without backward
Check config/flags that select the backend; training runs should build the device as Autodiff<YourBackend>

Example fix

// before
let device = CpuDevice::default();
device.gradient_checkpointing(); // panics: not an autodiff device
// after
let device = CpuDevice::default().autodiff();
device.gradient_checkpointing(); // ok: DispatchDevice::Autodiff

Defensive patterns

Strategy: type-guard

Validate before calling

// Ensure the device is autodiff-wrapped before enabling checkpointing
fn autodiff_first<B: AutodiffBackend>(device: &B::Device) { /* gradient_checkpointing is only valid on Autodiff devices */ }

Type guard

// match on the internal DispatchDevice kind before calling
fn is_autodiff_device(dev: &burn_tensor::Device) -> bool {
    matches!(dev, _ /* DispatchDevice::Autodiff(_) variant, per burn version */)
}

Try / catch

// Panic API; structure the code so checkpointing is only requested on autodiff devices:
let device = backend_device.autodiff();
device.gradient_checkpointing();

Prevention

Always create training devices as Autodiff<Backend> (device.autodiff()) before any training-specific configuration
Keep device construction in one factory function so autodiff wrapping can't be skipped
Only call gradient_checkpointing in training code paths gated on AutodiffBackend
Skip checkpointing entirely for inference runs

When it happens

Trigger: Calling gradient_checkpointing() directly on a device created from a non-autodiff backend (e.g. Wgpu/Cpu without Autodiff wrapper); calling it before Device::autodiff in the setup code; wrapping order mistakes where the autodiff layer was applied to a different device instance than the one checkpointing is requested on.

Common situations: Enabling gradient checkpointing to fit a large model in memory during training but forgetting the .autodiff() wrapper; copying training setup code between a script that used Autodiff<Backend> and one that uses a plain backend; refactors that moved autodiff wrapping into a config so the runtime device is no longer an autodiff variant.

Related errors

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