burn/burn-wgpu

Burn WGPU Backend

Burn WGPU backend

This crate provides a WGPU backend for Burn using the wgpu.

The backend supports Vulkan, Metal, DirectX11/12, OpenGL, WebGPU.

Usage Example

#[cfg(feature = "wgpu")]
mod wgpu {
    use burn_autodiff::Autodiff;
    use burn_wgpu::{AutoGraphicsApi, Wgpu, WgpuDevice};
    use mnist::training;

    pub fn run() {
        let device = WgpuDevice::default();
        training::run::<Autodiff<Wgpu<AutoGraphicsApi, f32, i32>>>(device);
    }
}

Configuration

You can set BURN_WGPU_MAX_TASKS to a positive integer that determines how many computing tasks are submitted in batches to the graphics API.