diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index cd6ae14..bd83c3a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -125,3 +125,38 @@ jobs:
 
       - name: Run ResNet-18 tests
         run: .venv/bin/python tests/test_resnet.py --bench
+
+  wgpu-tests:
+    name: wgpu tests (software Vulkan)
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Rust
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache cargo
+        uses: Swatinem/rust-cache@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      # no GPU on hosted runners: lavapipe is Mesa's CPU implementation of Vulkan,
+      # enough to check the GPU code path end to end (not its speed)
+      - name: Install lavapipe
+        run: sudo apt-get update && sudo apt-get install -y mesa-vulkan-drivers
+
+      - name: Create virtualenv
+        run: python -m venv .venv
+
+      - name: Install Python deps
+        run: .venv/bin/pip install maturin numpy onnx onnxruntime
+
+      - name: Build extension
+        run: source .venv/bin/activate && maturin develop --release
+
+      - name: Run op edge-case tests on wgpu
+        run: BURN_BACKEND=wgpu .venv/bin/python tests/test_ops.py
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63a5782..9251566 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,13 +7,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Performance
-- Static weights that onnx-ir passes as anonymous (unnamed) constants — including all
-  conv weights — are now pre-loaded once instead of converted on every forward pass.
-- BatchNormalization directly after a Conv2d is folded into the conv's weight and bias
-  at load time and removed from the graph. ResNet-18 (Model Zoo v1) at batch 8: 279 → 103 ms.
+### Added
+- GPU inference: `load_onnx(path, backend="wgpu")`, compiled to Metal on macOS and
+  Vulkan on Linux / Windows. ResNet-18 on an Apple M-series GPU: 6.1 ms at batch 1,
+  35.6 ms at batch 8 (~3x faster than ONNX Runtime on CPU). `available_backends()`
+  lists the backends in the build.
+- CI job running the op tests on wgpu through lavapipe (software Vulkan).
+- ReduceMean, ReduceSum, ReduceMax, ReduceMin (needed by the torch dynamo exporter's
+  global average pooling).
+- `tests/test_resnet.py`: ResNet-18 end to end vs ONNX Runtime on Model Zoo and
+  torchvision (legacy + dynamo) exports, with optional `--bench`; run in CI.
+- `tests/test_ops.py`: per-op edge-case tests against ONNX Runtime, run in CI.
+- ONNX interpreter support for Conv2d, fused BatchNormalization, MaxPool2d,
+  AveragePool2d, and GlobalAveragePool.
+- GitHub Actions CI: `cargo build`, `cargo fmt` / `cargo clippy` (non-blocking),
+  and Python tests via `maturin develop`.
+- Branch protection on `main`: pull requests required, `cargo build` and
+  `python tests (maturin)` checks required to merge.
+- ONNX interpreter with Gemm, Linear, Relu, Sigmoid, Tanh, Gelu, Softmax,
+  LogSoftmax, Add, Sub, Mul, Div, Reshape, Flatten, and Transpose ops.
+- numpy ↔ Burn tensor bridge (`roundtrip`), one copy in each direction.
+- Initial maturin/pyo3 project scaffold, `burn-flex` + `onnx-ir` dependencies.
 
 ### Changed
+- All ops now run on burn-dispatch's `Dispatch` backend; the device is chosen per model.
+  The unused `burn` umbrella dependency was replaced by `burn-dispatch`.
+- Backend failures that happen on a worker thread (e.g. no GPU adapter) now report the
+  worker's panic message instead of a downstream `RecvError`, and model loading is
+  guarded too, so they surface as `RuntimeError`.
 - Errors now surface as normal Python exceptions instead of `PanicException`:
   - `TypeError` for non-array / non-float32 inputs (with a `.astype(np.float32)` hint).
   - `ValueError` for wrong input count, rank, or a static dim that doesn't match the model.
@@ -22,7 +43,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Rust panic output on stderr is suppressed (the message is in the exception);
   set `RUST_BACKTRACE=1` to get it back.
 
+### Performance
+- Static weights that onnx-ir passes as anonymous (unnamed) constants — including all
+  conv weights — are now pre-loaded once instead of converted on every forward pass.
+- BatchNormalization directly after a Conv2d is folded into the conv's weight and bias
+  at load time and removed from the graph. ResNet-18 (Model Zoo v1) at batch 8: 279 → 103 ms.
+
 ### Fixed
+- Broadcasting between operands of different rank (`[N, C] + [C]`, Gemm's rank-1 `C`,
+  `[N, C, H, W] * [C, 1, 1]`) is now done numpy-style. It only worked by accident on
+  Flex; GPU backends rejected it.
+- Gemm `alpha`/`beta` and BatchNorm `epsilon` use scalar ops instead of 1-element tensors.
 - BatchNormalization on non-4D inputs (e.g. `[N, C]` after a Linear) no longer
   broadcasts to a 4D output; the pre-fused fast path is only used for NCHW.
 - Fortran-ordered inputs (e.g. `x.T`) were read in memory order, silently scrambling data.
@@ -36,21 +67,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Unsupported ops now fail at `load_onnx` with a list of the missing ops, instead of
   being skipped with a warning and failing (or silently misbehaving) at inference time.
 
-### Added
-- ReduceMean, ReduceSum, ReduceMax, ReduceMin (needed by the torch dynamo exporter's
-  global average pooling).
-- `tests/test_resnet.py`: ResNet-18 end to end vs ONNX Runtime on Model Zoo and
-  torchvision (legacy + dynamo) exports, with optional `--bench`; run in CI.
-- `tests/test_ops.py`: per-op edge-case tests against ONNX Runtime, run in CI.
-- ONNX interpreter support for Conv2d, fused BatchNormalization, MaxPool2d,
-  AveragePool2d, and GlobalAveragePool.
-- GitHub Actions CI: `cargo build`, `cargo fmt` / `cargo clippy` (non-blocking),
-  and Python tests via `maturin develop`.
-- Branch protection on `main`: pull requests required, `cargo build` and
-  `python tests (maturin)` checks required to merge.
-- ONNX interpreter with Gemm, Linear, Relu, Sigmoid, Tanh, Gelu, Softmax,
-  LogSoftmax, Add, Sub, Mul, Div, Reshape, Flatten, and Transpose ops.
-- Zero-copy numpy ↔ Burn tensor bridge (`roundtrip`).
-- Initial maturin/pyo3 project scaffold, `burn-flex` + `onnx-ir` dependencies.
-
 [Unreleased]: https://github.com/asjad2401/burn-python/commits/main
diff --git a/Cargo.toml b/Cargo.toml
index 3766d52..145fe9e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,14 +10,26 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3   = { version = "0.28", features = ["extension-module"] }
 numpy  = "0.28"
-burn   = { version = "0.21", features = ["flex"] }
 burn-flex    = "0.21"
+# listed directly so its default features (autotune, fusion) are on
+burn-wgpu    = "0.21"
 burn-backend = "0.21"
 onnx-ir      = "0.21"
 cubecl-common = "0.10"
 bytes     = "1"
 bytemuck  = { version = "1", features = ["extern_crate_alloc"] }
 
+# All ops go through burn-dispatch: Flex on CPU, wgpu on GPU. The GPU path is compiled
+# to each platform's native shading language — Metal (MSL) on macOS, Vulkan (SPIR-V)
+# elsewhere — which is several times faster than the portable WGSL path (no access to
+# matrix hardware). burn-dispatch allows only one wgpu flavor per build, hence per target.
+# wgpu loads the GPU driver at runtime, so CPU-only machines still work with "flex".
+[target.'cfg(target_os = "macos")'.dependencies]
+burn-dispatch = { version = "0.21", default-features = false, features = ["std", "flex", "metal"] }
+
+[target.'cfg(not(target_os = "macos"))'.dependencies]
+burn-dispatch = { version = "0.21", default-features = false, features = ["std", "flex", "vulkan"] }
+
 [profile.release]
 lto = true
 codegen-units = 1
diff --git a/README.md b/README.md
index 666201e..0108f90 100644
--- a/README.md
+++ b/README.md
@@ -15,6 +15,21 @@ x = np.random.randn(1, 3, 224, 224).astype(np.float32)
 output = model([x])[0]
 ```
 
+## Backends
+
+```python
+burn.available_backends()                          # ['flex', 'wgpu']
+model = burn.load_onnx("model.onnx", backend="wgpu")  # run on the GPU
+```
+
+| backend | runs on | notes |
+|---------|---------|-------|
+| `"flex"` (default) | CPU | pure-Rust [burn-flex](https://github.com/tracel-ai/burn) backend |
+| `"wgpu"` | GPU | Metal on macOS, Vulkan on Linux / Windows; first call per input shape compiles and autotunes kernels |
+
+Both use Burn's dispatch backend, so the same interpreter runs on either. If no usable
+GPU is found, `load_onnx(..., backend="wgpu")` raises a `RuntimeError` naming the problem.
+
 ## Status
 
 Early development. The numpy ↔ Burn tensor bridge is done, and the ONNX interpreter
@@ -33,15 +48,16 @@ ResNet-18 runs end to end and matches ONNX Runtime (max logit diff < 1e-5, same
 for the ONNX Model Zoo exports (`resnet18-v1-7`, `resnet18-v2-7`) and torchvision's
 `resnet18` exported with both `torch.onnx` exporters (legacy and dynamo).
 
-CPU latency vs ONNX Runtime (torchvision export, burn `flex` backend, median of 10 runs).
-Performance varies a lot by platform:
+Latency vs ONNX Runtime on CPU (torchvision export, median of 15 runs after warm-up).
+Numbers vary a lot by platform:
 
-| platform | batch | burn-python | ONNX Runtime | burn / ORT |
-|----------|------:|------------:|-------------:|-----------:|
-| Apple M-series (10 cores) | 1 | 22.6 ms  | 14.7 ms  | 1.5× slower |
-| Apple M-series (10 cores) | 8 | 105.0 ms | 118.5 ms | 0.9× (faster) |
-| Linux x86 (GitHub Actions runner) | 1 | 58.0 ms  | 18.7 ms  | 3.1× slower |
-| Linux x86 (GitHub Actions runner) | 8 | 335.7 ms | 145.9 ms | 2.3× slower |
+| platform | backend | batch 1 | batch 8 |
+|----------|---------|--------:|--------:|
+| Apple M-series (10 cores) | burn-python `wgpu` (Metal GPU) | **6.1 ms** | **35.6 ms** |
+| Apple M-series (10 cores) | burn-python `flex` (CPU) | 22.4 ms | 135.7 ms |
+| Apple M-series (10 cores) | ONNX Runtime (CPU) | 14.9 ms | 123.0 ms |
+| Linux x86 (GitHub Actions runner) | burn-python `flex` (CPU) | 58.0 ms | 335.7 ms |
+| Linux x86 (GitHub Actions runner) | ONNX Runtime (CPU) | 18.7 ms | 145.9 ms |
 
 Run `python tests/test_resnet.py --bench` to measure on your machine.
 
@@ -62,6 +78,8 @@ python tests/test_bridge.py       # numpy <-> Burn tensor bridge
 python tests/compare_ort.py       # correctness + perf vs ONNX Runtime
 python tests/test_ops.py          # op edge cases (padding, pooling, reshape) vs ORT
 python tests/test_resnet.py       # ResNet-18 end to end vs ORT (add --bench for timings)
+
+BURN_BACKEND=wgpu python tests/test_ops.py     # same tests on the GPU backend
 ```
 
 ## License
diff --git a/burn_python/__init__.py b/burn_python/__init__.py
index 6175338..75d59da 100644
--- a/burn_python/__init__.py
+++ b/burn_python/__init__.py
@@ -1,5 +1,5 @@
 import numpy as np
-from ._burn_python import __version__, OnnxModel
+from ._burn_python import __version__, OnnxModel, available_backends
 from . import _burn_python as _ext
 
 def roundtrip(arr: np.ndarray) -> np.ndarray:
@@ -9,9 +9,11 @@ def roundtrip(arr: np.ndarray) -> np.ndarray:
     return _ext.roundtrip(arr)
 
 def load_onnx(path: str, backend: str = "flex") -> OnnxModel:
-    """Load an ONNX model for inference."""
-    if backend != "flex":
-        raise ValueError(f"unsupported backend '{backend}' — only 'flex' is available right now")
-    return _ext.load_onnx(path)
+    """Load an ONNX model for inference.
 
-__all__ = ["__version__", "load_onnx", "roundtrip", "OnnxModel"]
+    backend: "flex" (CPU, default) or "wgpu" (GPU via Metal / Vulkan / DX12).
+    See available_backends() for what this build supports.
+    """
+    return _ext.load_onnx(str(path), backend)
+
+__all__ = ["__version__", "load_onnx", "available_backends", "roundtrip", "OnnxModel"]
diff --git a/src/interpreter/context.rs b/src/interpreter/context.rs
index 41435cd..2eddb3a 100644
--- a/src/interpreter/context.rs
+++ b/src/interpreter/context.rs
@@ -6,7 +6,7 @@ use std::collections::HashMap;
 use burn_backend::backend::ops::FloatTensorOps;
 use onnx_ir::ir::{Argument, ValueSource};
 
-use crate::tensor::{B, FloatPrim, default_device};
+use crate::tensor::{B, Device, FloatPrim};
 
 /// Key a static/constant argument is pre-loaded under in the weights map.
 /// onnx-ir often hands initializers over as anonymous statics (name ""), so those
@@ -25,13 +25,15 @@ pub type OpResult = Result<(), String>;
 pub struct ExecutionContext<'w> {
     tensors: HashMap<String, FloatPrim>,
     weights: &'w HashMap<String, FloatPrim>,
+    device: &'w Device,
 }
 
 impl<'w> ExecutionContext<'w> {
-    pub fn new(weights: &'w HashMap<String, FloatPrim>) -> Self {
+    pub fn new(weights: &'w HashMap<String, FloatPrim>, device: &'w Device) -> Self {
         Self {
             tensors: HashMap::new(),
             weights,
+            device,
         }
     }
 
@@ -56,10 +58,7 @@ impl<'w> ExecutionContext<'w> {
             ValueSource::Static(_) => weight_key(arg)
                 .and_then(|key| self.weights.get(&key).cloned())
                 // not pre-loaded (e.g. non-float data) — convert on the fly
-                .or_else(|| {
-                    arg.value()
-                        .map(|d| B::float_from_data(d, &default_device()))
-                }),
+                .or_else(|| arg.value().map(|d| B::float_from_data(d, self.device))),
             ValueSource::Optional => None,
         }
     }
diff --git a/src/interpreter/mod.rs b/src/interpreter/mod.rs
index 6ee9675..9756d72 100644
--- a/src/interpreter/mod.rs
+++ b/src/interpreter/mod.rs
@@ -5,7 +5,6 @@ mod context;
 pub mod ops;
 
 use std::collections::HashMap;
-use std::panic::{AssertUnwindSafe, catch_unwind};
 
 use burn_backend::{DType, TensorData, TensorMetadata, backend::ops::FloatTensorOps};
 use numpy::{
@@ -21,10 +20,10 @@ use pyo3::{
     prelude::*,
 };
 
-use crate::tensor::{B, FloatPrim, default_device, flex_to_numpy, numpy_to_flex};
+use crate::tensor::{B, Device, FloatPrim, device_for, numpy_to_tensor, tensor_to_numpy};
 use context::{ExecutionContext, OpResult, weight_key};
 
-fn load_weight(arg: &Argument) -> Option<(String, FloatPrim)> {
+fn load_weight(arg: &Argument, device: &Device) -> Option<(String, FloatPrim)> {
     match arg.value_source {
         ValueSource::Static(_) | ValueSource::Constant => {
             let key = weight_key(arg)?;
@@ -32,7 +31,7 @@ fn load_weight(arg: &Argument) -> Option<(String, FloatPrim)> {
             if !matches!(data.dtype, DType::F32 | DType::F16 | DType::BF16) {
                 return None;
             }
-            Some((key, B::float_from_data(data, &default_device())))
+            Some((key, B::float_from_data(data, device)))
         }
         _ => None,
     }
@@ -61,8 +60,7 @@ fn precompute_bn(
     let c = gamma.shape().iter().next().copied().unwrap_or(1);
 
     // scale = gamma / sqrt(var + eps),  shaped [1, C, 1, 1]
-    let eps_t = B::float_from_data(TensorData::from([eps]), &default_device());
-    let scale_flat = B::float_div(gamma, B::float_sqrt(B::float_add(var, eps_t)));
+    let scale_flat = B::float_div(gamma, B::float_sqrt(B::float_add_scalar(var, eps.into())));
     // offset = beta - mean * scale,  shaped [1, C, 1, 1]
     let offset_flat = B::float_sub(beta, B::float_mul(mean, scale_flat.clone()));
 
@@ -124,7 +122,7 @@ fn fold_conv_bn(
                 Some(b) => b,
                 None => continue, // dynamic bias: leave unfolded
             },
-            None => B::float_from_data(TensorData::zeros::<f32, _>([c_out]), &default_device()),
+            None => B::float_from_data(TensorData::zeros::<f32, _>([c_out]), &B::float_device(&w)),
         };
 
         let w = B::float_mul(
@@ -237,14 +235,7 @@ fn as_f32_array<'py>(
 
 /// Run one node, turning both op errors and backend panics into a message.
 fn run_node(node: &Node, ctx: &mut ExecutionContext) -> OpResult {
-    catch_unwind(AssertUnwindSafe(|| dispatch(node, ctx))).unwrap_or_else(|payload| {
-        let msg = payload
-            .downcast_ref::<String>()
-            .map(String::as_str)
-            .or_else(|| payload.downcast_ref::<&str>().copied())
-            .unwrap_or("unknown panic");
-        Err(msg.to_string())
-    })
+    crate::panic::catch(|| dispatch(node, ctx))?
 }
 
 #[pyclass]
@@ -253,6 +244,8 @@ pub struct OnnxModel {
     weights: HashMap<String, FloatPrim>,
     inputs: Vec<InputSpec>,
     output_names: Vec<String>,
+    backend: String,
+    device: Device,
 }
 
 #[pymethods]
@@ -270,11 +263,11 @@ impl OnnxModel {
             )));
         }
 
-        let mut ctx = ExecutionContext::new(&self.weights);
+        let mut ctx = ExecutionContext::new(&self.weights, &self.device);
         for (spec, obj) in self.inputs.iter().zip(&inputs) {
             let arr = as_f32_array(obj, &spec.name)?;
             spec.check_shape(arr.shape())?;
-            ctx.insert(spec.name.clone(), numpy_to_flex(&arr));
+            ctx.insert(spec.name.clone(), numpy_to_tensor(&arr, &self.device));
         }
 
         for node in &self.nodes {
@@ -293,7 +286,7 @@ impl OnnxModel {
                 let t = ctx.get(name).ok_or_else(|| {
                     PyRuntimeError::new_err(format!("model output '{name}' was never computed"))
                 })?;
-                flex_to_numpy(py, t)
+                tensor_to_numpy(py, t)
             })
             .collect()
     }
@@ -301,15 +294,26 @@ impl OnnxModel {
     fn __repr__(&self) -> String {
         let input_names: Vec<&str> = self.inputs.iter().map(|i| i.name.as_str()).collect();
         format!(
-            "OnnxModel(inputs={:?}, outputs={:?}, nodes={})",
+            "OnnxModel(inputs={:?}, outputs={:?}, nodes={}, backend={:?})",
             input_names,
             self.output_names,
-            self.nodes.len()
+            self.nodes.len(),
+            self.backend
         )
     }
 }
 
-pub fn load_onnx(path: &str) -> Result<OnnxModel, String> {
+pub fn load_onnx(path: &str, backend: &str) -> Result<OnnxModel, String> {
+    // backend init (e.g. no usable GPU) fails by panicking, often on a worker thread
+    crate::panic::catch(|| load_onnx_on(path, backend)).unwrap_or_else(|msg| {
+        Err(format!(
+            "failed to load model on backend '{backend}': {msg}"
+        ))
+    })
+}
+
+fn load_onnx_on(path: &str, backend: &str) -> Result<OnnxModel, String> {
+    let device = device_for(backend)?;
     let graph = OnnxGraphBuilder::new()
         .parse_file(path)
         .map_err(|e| e.to_string())?;
@@ -335,7 +339,7 @@ pub fn load_onnx(path: &str) -> Result<OnnxModel, String> {
     // pass 1: load all named static weights
     for node in &graph.nodes {
         for arg in node.inputs() {
-            if let Some((name, tensor)) = load_weight(arg) {
+            if let Some((name, tensor)) = load_weight(arg, &device) {
                 weights.insert(name, tensor);
             }
         }
@@ -360,6 +364,8 @@ pub fn load_onnx(path: &str) -> Result<OnnxModel, String> {
         output_names,
         nodes,
         weights,
+        backend: backend.to_string(),
+        device,
     })
 }
 
diff --git a/src/interpreter/ops/broadcast.rs b/src/interpreter/ops/broadcast.rs
new file mode 100644
index 0000000..7ade6aa
--- /dev/null
+++ b/src/interpreter/ops/broadcast.rs
@@ -0,0 +1,24 @@
+// ONNX uses numpy-style broadcasting, where operands may differ in rank
+// ([N, C] + [C], [N, C, H, W] * [C, 1, 1]). Burn broadcasts size-1 dims but needs
+// equal ranks, so the lower-rank operand gets leading 1-dims first.
+
+use burn_backend::{TensorMetadata, backend::ops::FloatTensorOps};
+
+use crate::tensor::{B, FloatPrim};
+
+fn unsqueeze_to(t: FloatPrim, rank: usize) -> FloatPrim {
+    let shape: Vec<usize> = t.shape().iter().copied().collect();
+    if shape.len() >= rank {
+        return t;
+    }
+    let padded: Vec<usize> = std::iter::repeat_n(1, rank - shape.len())
+        .chain(shape)
+        .collect();
+    B::float_reshape(t, padded.into())
+}
+
+/// Bring two operands to the same rank for an elementwise op.
+pub fn align(a: FloatPrim, b: FloatPrim) -> (FloatPrim, FloatPrim) {
+    let rank = a.shape().num_dims().max(b.shape().num_dims());
+    (unsqueeze_to(a, rank), unsqueeze_to(b, rank))
+}
diff --git a/src/interpreter/ops/conv.rs b/src/interpreter/ops/conv.rs
index 2999b0b..12591e9 100644
--- a/src/interpreter/ops/conv.rs
+++ b/src/interpreter/ops/conv.rs
@@ -1,8 +1,8 @@
 use super::super::context::{ExecutionContext, OpResult};
 use super::padding::{pad2d, resolve_pads, symmetric};
-use crate::tensor::{B, default_device};
+use crate::tensor::B;
 use burn_backend::{
-    TensorData, TensorMetadata,
+    TensorMetadata,
     backend::ops::ConvOptions,
     backend::ops::{FloatTensorOps, ModuleOps},
 };
@@ -78,8 +78,7 @@ pub fn batch_norm_fused(node: &BatchNormalizationNode, ctx: &mut ExecutionContex
             .chain(std::iter::repeat_n(1, rank - 2))
             .collect();
         let bcast = |t| B::float_reshape(t, bshape.clone().into());
-        let eps_t = B::float_from_data(TensorData::from([eps]), &default_device());
-        let denom = B::float_sqrt(B::float_add(bcast(var), eps_t));
+        let denom = B::float_sqrt(B::float_add_scalar(bcast(var), eps.into()));
         let y = B::float_div(B::float_sub(x, bcast(mean)), denom);
         let y = B::float_mul(y, bcast(gamma));
         let y = B::float_add(y, bcast(beta));
diff --git a/src/interpreter/ops/elementwise.rs b/src/interpreter/ops/elementwise.rs
index f15cddb..8253713 100644
--- a/src/interpreter/ops/elementwise.rs
+++ b/src/interpreter/ops/elementwise.rs
@@ -1,4 +1,5 @@
 use super::super::context::{ExecutionContext, OpResult};
+use super::broadcast::align;
 use crate::tensor::B;
 use burn_backend::backend::ops::FloatTensorOps;
 use onnx_ir::arithmetic::{AddNode, DivNode, MulNode, SubNode};
@@ -6,6 +7,7 @@ use onnx_ir::arithmetic::{AddNode, DivNode, MulNode, SubNode};
 pub fn add(node: &AddNode, ctx: &mut ExecutionContext) -> OpResult {
     let a = ctx.require(&node.inputs[0])?;
     let b = ctx.require(&node.inputs[1])?;
+    let (a, b) = align(a, b);
     ctx.insert(node.outputs[0].name.clone(), B::float_add(a, b));
     Ok(())
 }
@@ -13,6 +15,7 @@ pub fn add(node: &AddNode, ctx: &mut ExecutionContext) -> OpResult {
 pub fn sub(node: &SubNode, ctx: &mut ExecutionContext) -> OpResult {
     let a = ctx.require(&node.inputs[0])?;
     let b = ctx.require(&node.inputs[1])?;
+    let (a, b) = align(a, b);
     ctx.insert(node.outputs[0].name.clone(), B::float_sub(a, b));
     Ok(())
 }
@@ -20,6 +23,7 @@ pub fn sub(node: &SubNode, ctx: &mut ExecutionContext) -> OpResult {
 pub fn mul(node: &MulNode, ctx: &mut ExecutionContext) -> OpResult {
     let a = ctx.require(&node.inputs[0])?;
     let b = ctx.require(&node.inputs[1])?;
+    let (a, b) = align(a, b);
     ctx.insert(node.outputs[0].name.clone(), B::float_mul(a, b));
     Ok(())
 }
@@ -27,6 +31,7 @@ pub fn mul(node: &MulNode, ctx: &mut ExecutionContext) -> OpResult {
 pub fn div(node: &DivNode, ctx: &mut ExecutionContext) -> OpResult {
     let a = ctx.require(&node.inputs[0])?;
     let b = ctx.require(&node.inputs[1])?;
+    let (a, b) = align(a, b);
     ctx.insert(node.outputs[0].name.clone(), B::float_div(a, b));
     Ok(())
 }
diff --git a/src/interpreter/ops/linear.rs b/src/interpreter/ops/linear.rs
index 73fb52b..419c7f3 100644
--- a/src/interpreter/ops/linear.rs
+++ b/src/interpreter/ops/linear.rs
@@ -1,5 +1,6 @@
 use super::super::context::{ExecutionContext, OpResult};
-use crate::tensor::{B, default_device};
+use super::broadcast::align;
+use crate::tensor::B;
 use burn_backend::backend::ops::FloatTensorOps;
 use onnx_ir::{gemm::GemmNode, linear::LinearNode};
 
@@ -19,7 +20,10 @@ pub fn linear(node: &LinearNode, ctx: &mut ExecutionContext) -> OpResult {
 
     // bias is optional (3rd input)
     let y = match node.inputs.get(2).and_then(|a| ctx.resolve(a)) {
-        Some(b) => B::float_add(y, b),
+        Some(b) => {
+            let (y, b) = align(y, b);
+            B::float_add(y, b)
+        }
         None => y,
     };
 
@@ -47,24 +51,17 @@ pub fn gemm(node: &GemmNode, ctx: &mut ExecutionContext) -> OpResult {
     let mut y = B::float_matmul(a, b);
 
     if node.config.alpha != 1.0 {
-        let alpha = B::float_from_data(
-            burn_backend::TensorData::from([node.config.alpha]),
-            &default_device(),
-        );
-        y = B::float_mul(y, alpha);
+        y = B::float_mul_scalar(y, node.config.alpha.into());
     }
 
     if let Some(c) = node.inputs.get(2).and_then(|a| ctx.resolve(a)) {
         let c = if node.config.beta != 1.0 {
-            let beta = B::float_from_data(
-                burn_backend::TensorData::from([node.config.beta]),
-                &default_device(),
-            );
-            B::float_mul(c, beta)
+            B::float_mul_scalar(c, node.config.beta.into())
         } else {
             c
         };
-        y = B::float_add(y, c);
+        let (y_aligned, c) = align(y, c);
+        y = B::float_add(y_aligned, c);
     }
 
     ctx.insert(node.outputs[0].name.clone(), y);
diff --git a/src/interpreter/ops/mod.rs b/src/interpreter/ops/mod.rs
index 0e8e60e..5ab3f7a 100644
--- a/src/interpreter/ops/mod.rs
+++ b/src/interpreter/ops/mod.rs
@@ -1,4 +1,5 @@
 pub mod activation;
+mod broadcast;
 pub mod conv;
 pub mod elementwise;
 pub mod linear;
diff --git a/src/interpreter/ops/padding.rs b/src/interpreter/ops/padding.rs
index c341417..2d69d49 100644
--- a/src/interpreter/ops/padding.rs
+++ b/src/interpreter/ops/padding.rs
@@ -7,7 +7,7 @@
 use burn_backend::{Slice, TensorData, TensorMetadata, backend::ops::FloatTensorOps};
 use onnx_ir::node::padding::{AutoPad, PaddingConfig2d};
 
-use crate::tensor::{B, FloatPrim, default_device};
+use crate::tensor::{B, FloatPrim};
 
 /// Per-side pads for an NCHW tensor: [top, left, bottom, right].
 pub type Pads2d = [usize; 4];
@@ -63,7 +63,7 @@ pub fn pad2d(x: FloatPrim, pads: Pads2d, value: f32) -> FloatPrim {
 
     let padded = B::float_from_data(
         TensorData::full(vec![n, c, h + top + bottom, w + left + right], value),
-        &default_device(),
+        &B::float_device(&x),
     );
     let range =
         |start: usize, len: usize| Slice::new(start as isize, Some((start + len) as isize), 1);
diff --git a/src/interpreter/ops/pool.rs b/src/interpreter/ops/pool.rs
index 06a190c..26c6b1c 100644
--- a/src/interpreter/ops/pool.rs
+++ b/src/interpreter/ops/pool.rs
@@ -1,6 +1,6 @@
 use super::super::context::{ExecutionContext, OpResult};
 use super::padding::{pad2d, resolve_pads, symmetric};
-use crate::tensor::{B, default_device};
+use crate::tensor::B;
 use burn_backend::{
     TensorData, TensorMetadata,
     backend::ops::{FloatTensorOps, ModuleOps},
@@ -69,7 +69,7 @@ pub fn avg_pool2d(node: &AveragePool2dNode, ctx: &mut ExecutionContext) -> OpRes
             let shape = x.shape();
             let ones = B::float_from_data(
                 TensorData::full(vec![1, 1, shape[2], shape[3]], 1.0f32),
-                &default_device(),
+                &B::float_device(&x),
             );
             let sum = pool(pad2d(x, pads, 0.0), [0, 0], true);
             let count = pool(pad2d(ones, pads, 0.0), [0, 0], true);
diff --git a/src/lib.rs b/src/lib.rs
index a6099ae..62ea702 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@ use numpy::{PyArrayDyn, PyReadonlyArrayDyn};
 use pyo3::prelude::*;
 
 mod interpreter;
+mod panic;
 mod tensor;
 
 use interpreter::OnnxModel;
@@ -12,26 +13,31 @@ fn roundtrip<'py>(
     py: Python<'py>,
     arr: PyReadonlyArrayDyn<'py, f32>,
 ) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
-    let prim = tensor::numpy_to_flex(&arr);
-    tensor::flex_to_numpy(py, prim)
+    let prim = tensor::numpy_to_tensor(&arr, &tensor::cpu_device());
+    tensor::tensor_to_numpy(py, prim)
 }
 
-/// Load an ONNX model from a file path.
+/// Load an ONNX model from a file path, to run on the given backend.
 #[pyfunction]
-fn load_onnx(path: &str) -> PyResult<OnnxModel> {
-    interpreter::load_onnx(path).map_err(pyo3::exceptions::PyRuntimeError::new_err)
+#[pyo3(signature = (path, backend = "flex"))]
+fn load_onnx(path: &str, backend: &str) -> PyResult<OnnxModel> {
+    tensor::device_for(backend).map_err(pyo3::exceptions::PyValueError::new_err)?;
+    interpreter::load_onnx(path, backend).map_err(pyo3::exceptions::PyRuntimeError::new_err)
+}
+
+/// Backends compiled into this build.
+#[pyfunction]
+fn available_backends() -> Vec<&'static str> {
+    tensor::BACKENDS.to_vec()
 }
 
 #[pymodule]
 fn _burn_python(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    // Panics inside ops are caught and re-raised as Python exceptions that carry the
-    // message, so the default hook's stderr dump is just noise. RUST_BACKTRACE keeps it.
-    if std::env::var_os("RUST_BACKTRACE").is_none() {
-        std::panic::set_hook(Box::new(|_| {}));
-    }
+    panic::install_hook();
     m.add("__version__", env!("CARGO_PKG_VERSION"))?;
     m.add_function(wrap_pyfunction!(roundtrip, m)?)?;
     m.add_function(wrap_pyfunction!(load_onnx, m)?)?;
+    m.add_function(wrap_pyfunction!(available_backends, m)?)?;
     m.add_class::<OnnxModel>()?;
     Ok(())
 }
diff --git a/src/panic.rs b/src/panic.rs
new file mode 100644
index 0000000..229287a
--- /dev/null
+++ b/src/panic.rs
@@ -0,0 +1,46 @@
+// Turning Rust panics into Python exceptions.
+//
+// Panics inside Burn are caught and re-raised with their message, so the default hook's
+// stderr dump is just noise (RUST_BACKTRACE keeps it). GPU backends fail on a worker
+// thread, and the panic we catch is only a downstream `RecvError`; the hook records the
+// first panic on any thread so that root cause can be reported instead.
+
+use std::panic::{self, AssertUnwindSafe};
+use std::sync::Mutex;
+
+static FIRST_PANIC: Mutex<Option<String>> = Mutex::new(None);
+
+pub fn install_hook() {
+    let default_hook = panic::take_hook();
+    let verbose = std::env::var_os("RUST_BACKTRACE").is_some();
+    panic::set_hook(Box::new(move |info| {
+        if let Ok(mut first) = FIRST_PANIC.lock() {
+            first.get_or_insert_with(|| payload_message(info.payload()).to_string());
+        }
+        if verbose {
+            default_hook(info);
+        }
+    }));
+}
+
+fn payload_message(payload: &(dyn std::any::Any + Send)) -> &str {
+    payload
+        .downcast_ref::<String>()
+        .map(String::as_str)
+        .or_else(|| payload.downcast_ref::<&str>().copied())
+        .unwrap_or("unknown panic")
+}
+
+/// Run `f`, turning a panic (on this or a backend worker thread) into its message.
+pub fn catch<T>(f: impl FnOnce() -> T) -> Result<T, String> {
+    if let Ok(mut first) = FIRST_PANIC.lock() {
+        *first = None;
+    }
+    panic::catch_unwind(AssertUnwindSafe(f)).map_err(|payload| {
+        FIRST_PANIC
+            .lock()
+            .ok()
+            .and_then(|mut first| first.take())
+            .unwrap_or_else(|| payload_message(&*payload).to_string())
+    })
+}
diff --git a/src/tensor.rs b/src/tensor.rs
index cdee990..b0b3757 100644
--- a/src/tensor.rs
+++ b/src/tensor.rs
@@ -1,23 +1,53 @@
 // numpy <-> Burn tensor bridge
 //
-// in:  numpy f32 -> TensorData (one copy) -> FlexTensor
-// out: FlexTensor -> TensorData (sync CPU read) -> numpy f32 (one copy)
+// in:  numpy f32 -> TensorData (one copy) -> backend tensor (upload, for GPU backends)
+// out: backend tensor -> TensorData (sync read) -> numpy f32 (one copy)
+//
+// All ops run on burn-dispatch's `Dispatch` backend, which forwards each op to the
+// backend selected by the tensor's device (Flex on CPU, wgpu on GPU, ...).
 
 use burn_backend::{TensorData, backend::ops::FloatTensorOps};
-use burn_flex::{Flex, FlexDevice, FlexTensor};
+use burn_dispatch::{Dispatch, DispatchDevice, DispatchTensor};
+use burn_flex::FlexDevice;
 use cubecl_common::reader::read_sync;
 use numpy::{PyArray1, PyArrayDyn, PyArrayMethods, PyReadonlyArrayDyn, PyUntypedArrayMethods};
 use pyo3::{exceptions::PyRuntimeError, prelude::*};
 
-pub type B = Flex;
-pub type FloatPrim = FlexTensor;
+pub type B = Dispatch;
+pub type FloatPrim = DispatchTensor;
+pub type Device = DispatchDevice;
+
+/// Backend names accepted by `load_onnx(..., backend=...)`.
+pub const BACKENDS: &[&str] = &["flex", "wgpu"];
+
+/// Map a user-facing backend name to a device.
+pub fn device_for(backend: &str) -> Result<Device, String> {
+    match backend {
+        "flex" | "cpu" => Ok(cpu_device()),
+        "wgpu" | "gpu" => Ok(gpu_device()),
+        other => Err(format!(
+            "unknown or unavailable backend '{other}' (available: {})",
+            BACKENDS.join(", ")
+        )),
+    }
+}
+
+/// Default GPU, through the wgpu flavor this platform is built with (see Cargo.toml).
+fn gpu_device() -> Device {
+    let device = burn_wgpu::WgpuDevice::DefaultDevice;
+    #[cfg(target_os = "macos")]
+    return DispatchDevice::Metal(device);
+    #[cfg(not(target_os = "macos"))]
+    return DispatchDevice::Vulkan(device);
+}
 
-pub fn default_device() -> FlexDevice {
-    FlexDevice
+/// Device of the default CPU backend.
+pub fn cpu_device() -> Device {
+    DispatchDevice::Flex(FlexDevice)
 }
 
-// numpy f32 ndarray -> FlexTensor (one copy at the Python boundary)
-pub fn numpy_to_flex(arr: &PyReadonlyArrayDyn<'_, f32>) -> FloatPrim {
+// numpy f32 ndarray -> backend tensor (one copy at the Python boundary)
+pub fn numpy_to_tensor(arr: &PyReadonlyArrayDyn<'_, f32>, device: &Device) -> FloatPrim {
     let shape: Vec<usize> = arr.shape().to_vec();
     // as_slice() also accepts Fortran order (e.g. x.T), which would be read in memory
     // order; anything not C-contiguous is gathered in logical order instead
@@ -27,17 +57,19 @@ pub fn numpy_to_flex(arr: &PyReadonlyArrayDyn<'_, f32>) -> FloatPrim {
     };
     let data = TensorData::new(values, shape);
 
-    B::float_from_data(data, &default_device())
+    B::float_from_data(data, device)
 }
 
-// FlexTensor -> numpy f32 ndarray (sync read + one copy out)
-pub fn flex_to_numpy<'py>(
+// backend tensor -> numpy f32 ndarray (sync read + one copy out)
+pub fn tensor_to_numpy<'py>(
     py: Python<'py>,
     prim: FloatPrim,
 ) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
-    // for CPU backends the future resolves immediately
+    // blocks until the device has finished (immediate for CPU backends)
     let data: TensorData = read_sync(B::float_into_data(prim))
         .map_err(|e| PyRuntimeError::new_err(format!("failed to read tensor: {e:?}")))?;
+    // a backend may compute in another float type; the numpy side is always f32
+    let data = data.convert::<f32>();
 
     let shape: Vec<usize> = data.shape.iter().copied().collect();
     let floats: &[f32] = bytemuck::cast_slice(data.as_bytes());
diff --git a/tests/test_ops.py b/tests/test_ops.py
index 3befd61..3c21db4 100644
--- a/tests/test_ops.py
+++ b/tests/test_ops.py
@@ -5,6 +5,7 @@
 """
 
 import os
+import re
 import sys
 import tempfile
 import numpy as np
@@ -13,6 +14,8 @@
 from onnx import helper, numpy_helper, TensorProto
 import burn_python as burn
 
+BACKEND = os.environ.get("BURN_BACKEND", "flex")  # e.g. BURN_BACKEND=wgpu
+
 PASS = "\033[92mPASS\033[0m"
 FAIL = "\033[91mFAIL\033[0m"
 
@@ -36,7 +39,7 @@ def save_model(name, node, x_shape, initializers=(), opset=17):
     )
     model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)])
     model.ir_version = 10  # older onnxruntime builds reject newer IR versions
-    path = os.path.join(tmpdir, f"{name}.onnx")
+    path = os.path.join(tmpdir, re.sub(r"[^\w-]+", "_", name) + ".onnx")
     onnx.save(model, path)
     return path
 
@@ -44,7 +47,7 @@ def compare(label, node, x_shape, initializers=(), opset=17):
     path = save_model(label.replace(" ", "_"), node, x_shape, initializers, opset)
     x = np.random.randn(*x_shape).astype(np.float32)
     try:
-        got = burn.load_onnx(path)([x])[0]
+        got = burn.load_onnx(path, backend=BACKEND)([x])[0]
     except BaseException as e:  # pyo3 panics aren't Exception subclasses
         check(f"{label}  ({type(e).__name__}: {e})", False)
         return
@@ -112,9 +115,9 @@ def compare_graph(label, nodes, x_shape, initializers, outputs=("y",), n_nodes=N
     path = graph_model(label.replace(" ", "_"), nodes, x_shape, initializers, outputs)
     x = np.random.randn(*x_shape).astype(np.float32)
     try:
-        model = burn.load_onnx(path)
+        model = burn.load_onnx(path, backend=BACKEND)
         if n_nodes is not None:  # e.g. BN folded away into the conv
-            check(f"{label}  runs {n_nodes} node(s)", f"nodes={n_nodes})" in repr(model))
+            check(f"{label}  runs {n_nodes} node(s)", f"nodes={n_nodes}," in repr(model))
         got = model([x])
     except BaseException as e:
         check(f"{label}  ({type(e).__name__}: {e})", False)
@@ -169,6 +172,21 @@ def reshape(shape):
 node, init = reshape([-1, 4])
 compare("reshape [-1, 4]", node, (2, 3, 4), init)
 
+# ── broadcasting across ranks (numpy-style, as ONNX allows) ──────────────────
+
+print("\n=== broadcasting ===")
+
+for op, x_shape, c_shape in [("Add", (4, 3), (3,)), ("Mul", (2, 3, 5, 4), (3, 1, 1)),
+                             ("Sub", (2, 3, 5, 4), (4,)), ("Div", (3,), (2, 3))]:
+    c = numpy_helper.from_array(np.random.rand(*c_shape).astype(np.float32) + 0.5, "c")
+    compare(f"{op} {list(x_shape)} with {list(c_shape)}",
+            helper.make_node(op, ["x", "c"], ["y"]), x_shape, [c])
+
+gemm_c = numpy_helper.from_array(np.random.randn(2).astype(np.float32), "C")
+compare("Gemm alpha/beta, rank-1 C",
+        helper.make_node("Gemm", ["x", "W2", "C"], ["y"], alpha=0.5, beta=2.0),
+        (3, 4), [numpy_helper.from_array(np.random.randn(4, 2).astype(np.float32), "W2"), gemm_c])
+
 # ── reductions ───────────────────────────────────────────────────────────────
 
 print("\n=== reductions ===")
@@ -191,7 +209,7 @@ def reduce_node(op, axes, keepdims):
 
 path = save_model("unsupported", helper.make_node("Exp", ["x"], ["y"]), (2, 3))
 try:
-    burn.load_onnx(path)
+    burn.load_onnx(path, backend=BACKEND)
     check("unsupported op raises at load", False)
 except RuntimeError as e:
     check(f"unsupported op raises at load  ({e})", "Exp" in str(e))
@@ -205,11 +223,14 @@ def raises(label, fn, exc, *needles):
     except BaseException as e:  # e.g. pyo3 PanicException
         check(f"{label}  (wrong exception {type(e).__name__}: {e})", False)
 
+raises("unknown backend", lambda: burn.load_onnx(path, backend="tpu"), ValueError, "tpu", "flex")
+
 # ── inference errors ─────────────────────────────────────────────────────────
 
 print("\n=== inference errors ===")
 
-relu = burn.load_onnx(save_model("relu", helper.make_node("Relu", ["x"], ["y"]), ("N", 3)))
+relu = burn.load_onnx(save_model("relu", helper.make_node("Relu", ["x"], ["y"]), ("N", 3)),
+                      backend=BACKEND)
 raises("float64 input", lambda: relu([np.zeros((2, 3))]), TypeError, "float32", "float64")
 raises("non-array input", lambda: relu([[1.0, 2.0, 3.0]]), TypeError, "numpy array")
 raises("wrong rank", lambda: relu([np.zeros((2, 3, 1), np.float32)]), ValueError, "[?, 3]")
@@ -220,9 +241,10 @@ def raises(label, fn, exc, *needles):
 # doesn't match W, declared as fully dynamic so the input check lets it through
 gemm_w = numpy_helper.from_array(np.ones((4, 2), np.float32), "W")
 gemm = burn.load_onnx(save_model(
-    "gemm_dyn", helper.make_node("Gemm", ["x", "W"], ["y"]), ("N", "K"), [gemm_w]))
+    "gemm_dyn", helper.make_node("Gemm", ["x", "W"], ["y"]), ("N", "K"), [gemm_w]),
+    backend=BACKEND)
 raises("backend panic -> RuntimeError with node context",
-       lambda: gemm([np.zeros((2, 3), np.float32)]), RuntimeError, "(Gemm)", "inner dimensions")
+       lambda: gemm([np.zeros((2, 3), np.float32)]), RuntimeError, "(Gemm)")
 
 # ── non-contiguous inputs ────────────────────────────────────────────────────
 
diff --git a/tests/test_resnet.py b/tests/test_resnet.py
index 146952b..111ce99 100644
--- a/tests/test_resnet.py
+++ b/tests/test_resnet.py
@@ -15,6 +15,8 @@
 import onnxruntime as ort
 import burn_python as burn
 
+BACKEND = os.environ.get("BURN_BACKEND", "flex")  # e.g. BURN_BACKEND=wgpu
+
 MODEL_DIR = os.path.join(os.path.dirname(__file__), "models")
 ZOO_URL = "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/{}.onnx"
 BATCHES = [1, 4]
@@ -84,7 +86,7 @@ def median_ms(fn, runs=10):
 
 for path in models:
     print(f"\n=== {os.path.basename(path)} ===")
-    model = burn.load_onnx(path)
+    model = burn.load_onnx(path, backend=BACKEND)
     sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
     input_name = sess.get_inputs()[0].name