diff --git a/Cargo.lock b/Cargo.lock
index 275740a..2b78f0a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -255,15 +255,6 @@ dependencies = [
  "bit-vec 0.6.3",
 ]
 
-[[package]]
-name = "bit-set"
-version = "0.9.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "34ddef2995421ab6a5c779542c81ee77c115206f4ad9d5a8e05f4ff49716a3dd"
-dependencies = [
- "bit-vec 0.9.1",
-]
-
 [[package]]
 name = "bit-set"
 version = "0.10.0"
@@ -376,44 +367,28 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d450b8a4aefca19ce1dd51ee6980eaf379b6d238ced4de37535e75f72be7a532"
 dependencies = [
- "burn-autodiff 0.22.0-pre.2",
+ "burn-autodiff",
  "burn-core",
  "burn-cuda",
- "burn-ir 0.22.0-pre.2",
- "burn-ndarray 0.22.0-pre.2",
+ "burn-ir",
+ "burn-ndarray",
  "burn-nn",
  "burn-optim",
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "burn-store",
  "burn-train",
  "burn-vision",
  "burn-wgpu",
 ]
 
-[[package]]
-name = "burn-autodiff"
-version = "0.21.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3b93a80e43bfab909399444b29ce3894ff51f7e879720bfc75b6007ae6a01c3d"
-dependencies = [
- "burn-backend 0.21.0",
- "burn-std 0.21.0",
- "derive-new",
- "hashbrown 0.16.1",
- "log",
- "num-traits",
- "portable-atomic",
- "spin",
-]
-
 [[package]]
 name = "burn-autodiff"
 version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7c8a55150fdff1d1e8e1aca98010fb10f443e5808cdf1bad6cf7d0b096f1d7fe"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
+ "burn-backend",
+ "burn-std",
  "derive-new",
  "hashbrown 0.16.1",
  "log",
@@ -424,36 +399,15 @@ dependencies = [
  "spin",
 ]
 
-[[package]]
-name = "burn-backend"
-version = "0.21.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "64491519ac0867f550fa9c726cd443e5db94608c629050986cf2614c8c5ad39f"
-dependencies = [
- "burn-std 0.21.0",
- "bytemuck",
- "cubecl 0.10.0",
- "derive-new",
- "enumset",
- "hashbrown 0.16.1",
- "num-traits",
- "portable-atomic-util",
- "rand 0.10.1",
- "rand_distr",
- "serde",
- "spin",
- "thiserror 2.0.18",
-]
-
 [[package]]
 name = "burn-backend"
 version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b8157cd9dfe164c4a3e00183a4fbaa26789209ecf13ba55906185b532a20f1c5"
 dependencies = [
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "bytemuck",
- "cubecl 0.11.0-pre.2",
+ "cubecl",
  "derive-new",
  "enumset",
  "hashbrown 0.16.1",
@@ -484,13 +438,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7861d5f03dfa489c05313cfd69bb02be59f4cece8cdc54731e0bb248abd2bb00"
 dependencies = [
  "ahash",
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-backend-extension",
  "burn-dataset",
  "burn-derive",
  "burn-dispatch",
  "burn-pack",
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "burn-tensor",
  "derive-new",
  "half",
@@ -512,10 +466,10 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "357ce26165dcaf27c8acfdff97170a6ce724db3c4076252fbb0669ae274fec53"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-cubecl",
  "burn-fusion",
- "cubecl 0.11.0-pre.2",
+ "cubecl",
 ]
 
 [[package]]
@@ -524,12 +478,12 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3a1db3d0cd0d081e2dd0b8dce41847bc88495cc3e9e89c47c391f8c6940263bd"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-cubecl-fusion",
  "burn-fusion",
- "burn-ir 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
- "cubecl 0.11.0-pre.2",
+ "burn-ir",
+ "burn-std",
+ "cubecl",
  "cubek",
  "derive-new",
  "futures-lite",
@@ -544,11 +498,11 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "40c80f796e4c88fd196022ca1f12a5d04ae28f175cb8be82ca364446eb7b0e7f"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-fusion",
- "burn-ir 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
- "cubecl 0.11.0-pre.2",
+ "burn-ir",
+ "burn-std",
+ "cubecl",
  "cubek",
  "derive-new",
  "hashbrown 0.16.1",
@@ -563,10 +517,10 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3d2f7d0d2dfa0221d87e095944807f38fb201582b9b9c62d72f47019a5284b11"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-cubecl",
  "burn-fusion",
- "cubecl 0.11.0-pre.2",
+ "cubecl",
 ]
 
 [[package]]
@@ -606,12 +560,12 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b5bf565da978d065420c0b750417e5bc6ee72e5a168e0b2b15fc300439a6d622"
 dependencies = [
- "burn-autodiff 0.22.0-pre.2",
- "burn-backend 0.22.0-pre.2",
+ "burn-autodiff",
+ "burn-backend",
  "burn-cpu",
  "burn-cuda",
  "burn-flex",
- "burn-ndarray 0.22.0-pre.2",
+ "burn-ndarray",
  "burn-remote",
  "burn-rocm",
  "burn-tch",
@@ -626,9 +580,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3af60c5cda3823448c723d6e3840223e806bb122a639689421a2ca1164decbaa"
 dependencies = [
  "aligned-vec",
- "burn-backend 0.22.0-pre.2",
- "burn-ir 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
+ "burn-backend",
+ "burn-ir",
+ "burn-std",
  "bytemuck",
  "gemm",
  "half",
@@ -647,9 +601,9 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0752c69048173820fd19cdbcb11f50e780dc4aaa2fee9f68ab75a6e0365b975e"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
- "burn-ir 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
+ "burn-backend",
+ "burn-ir",
+ "burn-std",
  "derive-new",
  "hashbrown 0.16.1",
  "log",
@@ -659,55 +613,17 @@ dependencies = [
  "tracing",
 ]
 
-[[package]]
-name = "burn-ir"
-version = "0.21.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a81a55c32f35c2265d9e15ba2e796013d9780d5a49f9e21ea0ac4d29609dbf13"
-dependencies = [
- "burn-backend 0.21.0",
- "hashbrown 0.16.1",
- "serde",
-]
-
 [[package]]
 name = "burn-ir"
 version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "26d4ac0a2fcd725d7d429589944f2b894bc1ff5da01d40598f46b9561fc8b717"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "hashbrown 0.16.1",
  "serde",
 ]
 
-[[package]]
-name = "burn-ndarray"
-version = "0.21.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dceb9692e292782bab7ad0259e8e8f5ee80ba2b0329a78f0ba589a26a9633dd6"
-dependencies = [
- "atomic_float",
- "burn-autodiff 0.21.0",
- "burn-backend 0.21.0",
- "burn-ir 0.21.0",
- "burn-std 0.21.0",
- "bytemuck",
- "const-random",
- "itertools",
- "libm",
- "macerator",
- "matrixmultiply",
- "ndarray 0.17.2",
- "num-traits",
- "paste",
- "portable-atomic",
- "portable-atomic-util",
- "rand 0.10.1",
- "rayon",
- "seq-macro",
-]
-
 [[package]]
 name = "burn-ndarray"
 version = "0.22.0-pre.2"
@@ -715,9 +631,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5e59e56039e0ddd460741b95ccc5cb275de365c3127927a767b16421f0b4e157"
 dependencies = [
  "atomic_float",
- "burn-backend 0.22.0-pre.2",
- "burn-ir 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
+ "burn-backend",
+ "burn-ir",
+ "burn-std",
  "const-random",
  "libm",
  "macerator",
@@ -762,7 +678,7 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "93325a6545a605b5079d77957970b26f920672c12555589b34ae67242e3f4717"
 dependencies = [
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "byteorder",
  "ciborium",
  "serde",
@@ -774,11 +690,11 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fa8db206a100d838eeb100b8525ff8e3051958ee822fcfecd736fe58eddaeeaa"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-fusion",
- "burn-ir 0.22.0-pre.2",
+ "burn-ir",
  "burn-router",
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "bytes",
  "futures-util",
  "gloo-timers",
@@ -797,10 +713,10 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8667b9c5003d82f2c0511a8a8aef0bddf68874f261d549653a3e7e0fabe37164"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-cubecl",
  "burn-fusion",
- "cubecl 0.11.0-pre.2",
+ "cubecl",
 ]
 
 [[package]]
@@ -809,33 +725,16 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "91cc6e085ed958e988840a4fd3c4ec684a0479dcd3089bcf1b55babd8735eec7"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-fusion",
- "burn-ir 0.22.0-pre.2",
- "burn-std 0.22.0-pre.2",
+ "burn-ir",
+ "burn-std",
  "hashbrown 0.16.1",
  "log",
  "serde",
  "spin",
 ]
 
-[[package]]
-name = "burn-std"
-version = "0.21.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "81c7eb60eb2c316854af5b214fd55f59fc7b0e25dfde7db671f09cc49f269048"
-dependencies = [
- "bytemuck",
- "bytes",
- "cubecl-common 0.10.0",
- "cubecl-zspace 0.10.0",
- "half",
- "num-traits",
- "serde",
- "smallvec",
- "spin",
-]
-
 [[package]]
 name = "burn-std"
 version = "0.22.0-pre.2"
@@ -845,9 +744,9 @@ dependencies = [
  "ahash",
  "bytemuck",
  "bytes",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl-common",
  "cubecl-environment",
- "cubecl-zspace 0.11.0-pre.2",
+ "cubecl-zspace",
  "data-encoding",
  "derive-new",
  "enumset",
@@ -888,7 +787,7 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5b5ff862775a74d15b7e11205368e9cc21e3d79da62be91271efdc90a62e737d"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "cc",
  "libc",
  "log",
@@ -902,9 +801,9 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "85e85e6beb63d3ccd18fa5c1f71e8e3806a7f940cec8769f522b9d8cd173c999"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-dispatch",
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "colored",
  "derive-new",
  "num-traits",
@@ -921,7 +820,7 @@ dependencies = [
  "burn-core",
  "burn-flex",
  "burn-optim",
- "burn-std 0.22.0-pre.2",
+ "burn-std",
  "derive-new",
  "log",
  "rand 0.10.1",
@@ -946,9 +845,9 @@ dependencies = [
  "burn-core",
  "burn-cubecl",
  "burn-fusion",
- "burn-ir 0.22.0-pre.2",
+ "burn-ir",
  "bytemuck",
- "cubecl 0.11.0-pre.2",
+ "cubecl",
  "derive-new",
  "half",
  "image",
@@ -965,10 +864,10 @@ version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "798e66d9efc792e6e646e16d6f531a2c6e503bbe92f2b017c84f2306fcc9e2d3"
 dependencies = [
- "burn-backend 0.22.0-pre.2",
+ "burn-backend",
  "burn-cubecl",
  "burn-fusion",
- "cubecl 0.11.0-pre.2",
+ "cubecl",
 ]
 
 [[package]]
@@ -1407,79 +1306,24 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
-[[package]]
-name = "cubecl"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fd203fef6e359472e4cb8f6c0ef3eca062815a1edd560eb6d6c1d5c1397838d9"
-dependencies = [
- "cubecl-core 0.10.0",
- "cubecl-cuda 0.10.0",
- "cubecl-ir 0.10.0",
- "cubecl-runtime 0.10.0",
- "cubecl-wgpu 0.10.0",
- "half",
-]
-
 [[package]]
 name = "cubecl"
 version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2029987fbb8d073eabc5ca09887846c2db6928fa81f90afc6b883ce819425884"
 dependencies = [
- "cubecl-core 0.11.0-pre.2",
+ "cubecl-core",
  "cubecl-cpu",
- "cubecl-cuda 0.11.0-pre.2",
+ "cubecl-cuda",
  "cubecl-environment",
  "cubecl-hip",
- "cubecl-ir 0.11.0-pre.2",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-ir",
+ "cubecl-runtime",
  "cubecl-std",
- "cubecl-wgpu 0.11.0-pre.2",
+ "cubecl-wgpu",
  "half",
 ]
 
-[[package]]
-name = "cubecl-common"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cc956c2dcc993f16f748d03bfdbd900f75cab4d469f4e5d8924f8ee128e861ad"
-dependencies = [
- "backtrace",
- "bytemuck",
- "bytes",
- "cfg-if",
- "cfg_aliases",
- "ciborium",
- "derive-new",
- "derive_more",
- "dirs",
- "embassy-futures",
- "embassy-time",
- "float4",
- "float8",
- "futures-lite",
- "half",
- "hashbrown 0.16.1",
- "log",
- "num-traits",
- "oneshot",
- "parking_lot",
- "portable-atomic",
- "portable-atomic-util",
- "rand 0.10.1",
- "sanitize-filename",
- "serde",
- "serde_bytes",
- "serde_json",
- "spin",
- "toml",
- "tynm",
- "wasm-bindgen-futures",
- "web-time",
- "xxhash-rust",
-]
-
 [[package]]
 name = "cubecl-common"
 version = "0.11.0-pre.2"
@@ -1505,33 +1349,6 @@ dependencies = [
  "xxhash-rust",
 ]
 
-[[package]]
-name = "cubecl-core"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a7522e7acf25d7848032c7b5270780f9f2abe84e13c7ed6345aa27ba70c312ca"
-dependencies = [
- "bitflags 2.11.1",
- "bytemuck",
- "cubecl-common 0.10.0",
- "cubecl-ir 0.10.0",
- "cubecl-macros 0.10.0",
- "cubecl-runtime 0.10.0",
- "cubecl-zspace 0.10.0",
- "derive-new",
- "derive_more",
- "enumset",
- "float-ord",
- "half",
- "hashbrown 0.16.1",
- "log",
- "num-traits",
- "paste",
- "serde",
- "serde_json",
- "variadics_please 1.1.0",
-]
-
 [[package]]
 name = "cubecl-core"
 version = "0.11.0-pre.2"
@@ -1540,12 +1357,12 @@ checksum = "6d6fd5fb5fbf719e1a86df673dddae9e5183b08f0744308b2558fcc0632f5fe9"
 dependencies = [
  "bitflags 2.11.1",
  "bytemuck",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl-common",
  "cubecl-environment",
- "cubecl-ir 0.11.0-pre.2",
- "cubecl-macros 0.11.0-pre.2",
- "cubecl-runtime 0.11.0-pre.2",
- "cubecl-zspace 0.11.0-pre.2",
+ "cubecl-ir",
+ "cubecl-macros",
+ "cubecl-runtime",
+ "cubecl-zspace",
  "derive-new",
  "derive_more",
  "enumset",
@@ -1560,24 +1377,7 @@ dependencies = [
  "serde",
  "serde_json",
  "tracing",
- "variadics_please 2.0.0",
-]
-
-[[package]]
-name = "cubecl-cpp"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "411af04828cbf4583ecd388579fb30cd7a644eec19a334bb8f0d9d08522e1458"
-dependencies = [
- "bytemuck",
- "cubecl-common 0.10.0",
- "cubecl-core 0.10.0",
- "cubecl-opt 0.10.0",
- "cubecl-runtime 0.10.0",
- "derive-new",
- "half",
- "itertools",
- "log",
+ "variadics_please",
 ]
 
 [[package]]
@@ -1587,11 +1387,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1f43128c27f790ff22012679945c2b8972296979c52e5c323b54be830069e6e9"
 dependencies = [
  "bytemuck",
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
  "cubecl-environment",
- "cubecl-opt 0.11.0-pre.2",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-opt",
+ "cubecl-runtime",
  "derive-new",
  "half",
  "itertools",
@@ -1606,11 +1406,11 @@ checksum = "8274cac78ff1528b075eb0240b576690c419e13c4621dc1c596ea0a6fce809ca"
 dependencies = [
  "bytemuck",
  "crossbeam-utils",
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
  "cubecl-environment",
- "cubecl-opt 0.11.0-pre.2",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-opt",
+ "cubecl-runtime",
  "cubecl-std",
  "derive-new",
  "half",
@@ -1626,24 +1426,6 @@ dependencies = [
  "winapi",
 ]
 
-[[package]]
-name = "cubecl-cuda"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b6b0a69ff45688d322ad8e92c8bf645167b9ca490fa8fa087fc6adac8c5e46be"
-dependencies = [
- "bytemuck",
- "cubecl-common 0.10.0",
- "cubecl-core 0.10.0",
- "cubecl-cpp 0.10.0",
- "cubecl-runtime 0.10.0",
- "cudarc",
- "derive-new",
- "half",
- "log",
- "serde",
-]
-
 [[package]]
 name = "cubecl-cuda"
 version = "0.11.0-pre.2"
@@ -1651,11 +1433,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3e51bbfb4834e79972d7248bb8f69a747e848415f6e06e01aa8529cc8a29ad69"
 dependencies = [
  "bytemuck",
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
- "cubecl-cpp 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
+ "cubecl-cpp",
  "cubecl-environment",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-runtime",
  "cudarc",
  "derive-new",
  "half",
@@ -1707,12 +1489,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7eaad6c93c5c48aa075457f03d8eeeca6f3ccb01f19a7da4940c664d48a14248"
 dependencies = [
  "bytemuck",
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
- "cubecl-cpp 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
+ "cubecl-cpp",
  "cubecl-environment",
  "cubecl-hip-sys",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-runtime",
  "derive-new",
  "half",
  "log",
@@ -1731,28 +1513,6 @@ dependencies = [
  "regex",
 ]
 
-[[package]]
-name = "cubecl-ir"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d28a897a40c8ee6c57a8b8d76d8823ba2e97f4c2b83db9338f17a0752132d486"
-dependencies = [
- "cubecl-common 0.10.0",
- "cubecl-macros-internal 0.10.0",
- "derive-new",
- "derive_more",
- "enumset",
- "float-ord",
- "fnv",
- "foldhash 0.2.0",
- "half",
- "hashbrown 0.16.1",
- "num-traits",
- "portable-atomic",
- "serde",
- "variadics_please 1.1.0",
-]
-
 [[package]]
 name = "cubecl-ir"
 version = "0.11.0-pre.2"
@@ -1760,9 +1520,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "17ca2b43b9cfb6c4013592fa023a9f9158e344074d9dbba040b45461cc0842bc"
 dependencies = [
  "bumpalo",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl-common",
  "cubecl-environment",
- "cubecl-macros-internal 0.11.0-pre.2",
+ "cubecl-macros-internal",
  "derive-new",
  "derive_more",
  "enumset",
@@ -1776,24 +1536,7 @@ dependencies = [
  "num-traits",
  "portable-atomic",
  "serde",
- "variadics_please 2.0.0",
-]
-
-[[package]]
-name = "cubecl-macros"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77aa6df563e8e6a0926d2e9eeacb968737940a0e1ae9be819f6a65a341006e12"
-dependencies = [
- "cubecl-common 0.10.0",
- "darling 0.23.0",
- "derive-new",
- "ident_case",
- "inflections",
- "prettyplease",
- "proc-macro2",
- "quote",
- "syn 2.0.117",
+ "variadics_please",
 ]
 
 [[package]]
@@ -1802,7 +1545,7 @@ version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f6b8fa7c11f7e311af88fc12883eedef8e0350541121fe8cc634db335c6e49ce"
 dependencies = [
- "cubecl-common 0.11.0-pre.2",
+ "cubecl-common",
  "darling 0.23.0",
  "derive-new",
  "derive_more",
@@ -1814,18 +1557,6 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "cubecl-macros-internal"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a515651a5a91e25d87f71f311f80a088e6cf815798b9922560a97636da6b8740"
-dependencies = [
- "darling 0.23.0",
- "proc-macro2",
- "quote",
- "syn 2.0.117",
-]
-
 [[package]]
 name = "cubecl-macros-internal"
 version = "0.11.0-pre.2"
@@ -1838,34 +1569,16 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "cubecl-opt"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0a599c6c3efefdcee8636994c90e58ef696c7efbed2d18b5b5ad8d5f29923abb"
-dependencies = [
- "cubecl-common 0.10.0",
- "cubecl-core 0.10.0",
- "cubecl-ir 0.10.0",
- "float-ord",
- "log",
- "num",
- "petgraph",
- "smallvec",
- "stable-vec",
- "type-map",
-]
-
 [[package]]
 name = "cubecl-opt"
 version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d0904a499f3ccbbfb24c2c8668d516f351c3e50c9178068be87d92bb3b8e1478"
 dependencies = [
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
  "cubecl-environment",
- "cubecl-ir 0.11.0-pre.2",
+ "cubecl-ir",
  "float-ord",
  "hashbrown 0.17.1",
  "log",
@@ -1876,36 +1589,6 @@ dependencies = [
  "type-map",
 ]
 
-[[package]]
-name = "cubecl-runtime"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b68491bf5b3e997ae36bdc4e63b4ccd6d2f0e86b3b596a5d7a48d2b9e92622a0"
-dependencies = [
- "ahash",
- "async-channel",
- "bytemuck",
- "cfg-if",
- "cfg_aliases",
- "cubecl-common 0.10.0",
- "cubecl-ir 0.10.0",
- "cubecl-zspace 0.10.0",
- "derive-new",
- "derive_more",
- "dirs",
- "enumset",
- "hashbrown 0.16.1",
- "log",
- "md5",
- "serde",
- "serde_json",
- "spin",
- "thiserror 2.0.18",
- "toml",
- "wasm-bindgen-futures",
- "web-time",
-]
-
 [[package]]
 name = "cubecl-runtime"
 version = "0.11.0-pre.2"
@@ -1916,10 +1599,10 @@ dependencies = [
  "bytemuck",
  "cfg-if",
  "cfg_aliases",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl-common",
  "cubecl-environment",
- "cubecl-ir 0.11.0-pre.2",
- "cubecl-zspace 0.11.0-pre.2",
+ "cubecl-ir",
+ "cubecl-zspace",
  "derive-new",
  "derive_more",
  "enumset",
@@ -1941,40 +1624,16 @@ version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bd029b15a150c3326b71a436f20b81973a2898ae9406a349c726918efdacf455"
 dependencies = [
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
  "cubecl-environment",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-runtime",
  "half",
  "num-traits",
  "paste",
  "serde",
  "spin",
- "variadics_please 2.0.0",
-]
-
-[[package]]
-name = "cubecl-wgpu"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4d3663c6e1a86187172b2cee495b6b533e7d91a2cb37e26f421649e315fe4c3a"
-dependencies = [
- "async-channel",
- "bytemuck",
- "cfg-if",
- "cfg_aliases",
- "cubecl-common 0.10.0",
- "cubecl-core 0.10.0",
- "cubecl-ir 0.10.0",
- "cubecl-runtime 0.10.0",
- "derive-new",
- "derive_more",
- "half",
- "hashbrown 0.16.1",
- "log",
- "sanitize-filename",
- "wasm-bindgen-futures",
- "wgpu 29.0.3",
+ "variadics_please",
 ]
 
 [[package]]
@@ -1986,11 +1645,11 @@ dependencies = [
  "bytemuck",
  "cfg-if",
  "cfg_aliases",
- "cubecl-common 0.11.0-pre.2",
- "cubecl-core 0.11.0-pre.2",
+ "cubecl-common",
+ "cubecl-core",
  "cubecl-environment",
- "cubecl-ir 0.11.0-pre.2",
- "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-ir",
+ "cubecl-runtime",
  "derive-new",
  "derive_more",
  "half",
@@ -2000,18 +1659,7 @@ dependencies = [
  "sanitize-filename",
  "tracing",
  "wasm-bindgen-futures",
- "wgpu 30.0.0",
-]
-
-[[package]]
-name = "cubecl-zspace"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "04054d95698afab785a4dda9f949e297831dd9e4cc6d036a93e6603f88e88b07"
-dependencies = [
- "derive-new",
- "serde",
- "smallvec",
+ "wgpu",
 ]
 
 [[package]]
@@ -2031,7 +1679,7 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7f51262bc23900f4dc70c0a41f8d56910c3c6971e1dc242ec687cf920eda55c2"
 dependencies = [
- "cubecl 0.11.0-pre.2",
+ "cubecl",
  "cubek-attention",
  "cubek-convolution",
  "cubek-fft",
@@ -2051,8 +1699,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "739d905b83b5e42854cdfb8201389c688636043acc36c882b71a7ac10735e65b"
 dependencies = [
  "bytemuck",
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "cubek-matmul",
  "cubek-std",
  "half",
@@ -2066,8 +1714,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b3c24e82dad0d0d0c6fcaf7d13376653529e91639ef401457e972f2e329fb018"
 dependencies = [
  "bytemuck",
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "cubek-matmul",
  "cubek-std",
  "derive-new",
@@ -2082,7 +1730,7 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ddf0f55951bfd0d0ceddb67fc14f34e7484ef6a1838b9522ad2bd36c5cba11c3"
 dependencies = [
- "cubecl 0.11.0-pre.2",
+ "cubecl",
 ]
 
 [[package]]
@@ -2091,8 +1739,8 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8ce68e9c94bb9274827aca056c2589b24ce5acf2abba6aa8840bc27841eaad99"
 dependencies = [
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "half",
  "serde",
  "thiserror 2.0.18",
@@ -2105,8 +1753,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "204917532e8f5bc4b440c640551d2c10d402bac7fd8062450d843752352518df"
 dependencies = [
  "bytemuck",
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "cubek-std",
  "cubek-tile",
  "half",
@@ -2119,8 +1767,8 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9e09ab3dd3a448b7672bc4ce4a880909896c733fb42e958c035d76dd6bbf291"
 dependencies = [
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "thiserror 2.0.18",
 ]
 
@@ -2130,8 +1778,8 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "76384f2666b4c552881e146af6fe543aecf422ecff4d9ffbefb75b91461cf0e5"
 dependencies = [
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "cubek-tile",
  "half",
  "serde",
@@ -2143,8 +1791,8 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eb20538cb2a0d548a5571846021911d4ee660f75b1c10f6237c22cce1a2d4f7e"
 dependencies = [
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "cubecl-environment",
  "half",
  "num-traits",
@@ -2158,7 +1806,7 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5f05c73af6d1b35f886661b44813985284d86f6ce6d3395c064f52c6e8d50f81"
 dependencies = [
- "cubecl 0.11.0-pre.2",
+ "cubecl",
  "cubek-std",
  "half",
  "num-traits",
@@ -2172,8 +1820,8 @@ version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "db70badd24b211310d9db3ccc280ccd148100f5f42a112c594fbc0e6b3c9a5dd"
 dependencies = [
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "half",
 ]
 
@@ -2184,8 +1832,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b21695c6662edd669148874aea0eb8f0b011c7ef5bf15efa242b8bdae5d71555"
 dependencies = [
  "bytemuck",
- "cubecl 0.11.0-pre.2",
- "cubecl-common 0.11.0-pre.2",
+ "cubecl",
+ "cubecl-common",
  "half",
  "serde",
 ]
@@ -2519,7 +2167,6 @@ name = "engine"
 version = "0.1.0"
 dependencies = [
  "burn",
- "burn-ndarray 0.21.0",
  "chess",
  "ctrlc",
  "rand 0.10.1",
@@ -3159,26 +2806,6 @@ dependencies = [
  "windows",
 ]
 
-[[package]]
-name = "gpu-descriptor"
-version = "0.3.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b89c83349105e3732062a895becfc71a8f921bb71ecbbdd8ff99263e3b53a0ca"
-dependencies = [
- "bitflags 2.11.1",
- "gpu-descriptor-types",
- "hashbrown 0.15.5",
-]
-
-[[package]]
-name = "gpu-descriptor-types"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fdf242682df893b86f33a73828fb09ca4b2d3bb6cc95249707fc684d27484b91"
-dependencies = [
- "bitflags 2.11.1",
-]
-
 [[package]]
 name = "h2"
 version = "0.4.14"
@@ -3285,12 +2912,6 @@ version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
 
-[[package]]
-name = "hexf-parse"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df"
-
 [[package]]
 name = "hmac"
 version = "0.12.1"
@@ -3940,10 +3561,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
 dependencies = [
  "autocfg",
- "num_cpus",
- "once_cell",
  "rawpointer",
- "thread-tree",
 ]
 
 [[package]]
@@ -4042,32 +3660,6 @@ version = "0.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bc0287524726960e07b119cebd01678f852f147742ae0d925e6a520dca956126"
 
-[[package]]
-name = "naga"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0dd91265cc2454558f659b3b4b9640f0ddb8cc6521277f166b8a8c181c898079"
-dependencies = [
- "arrayvec 0.7.6",
- "bit-set 0.9.1",
- "bitflags 2.11.1",
- "cfg-if",
- "cfg_aliases",
- "codespan-reporting",
- "half",
- "hashbrown 0.16.1",
- "hexf-parse",
- "indexmap",
- "libm",
- "log",
- "num-traits",
- "once_cell",
- "rustc-hash 1.1.0",
- "spirv",
- "thiserror 2.0.18",
- "unicode-ident",
-]
-
 [[package]]
 name = "naga"
 version = "30.0.0"
@@ -4149,7 +3741,6 @@ dependencies = [
  "portable-atomic",
  "portable-atomic-util",
  "rawpointer",
- "rayon",
 ]
 
 [[package]]
@@ -6175,15 +5766,6 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "thread-tree"
-version = "0.3.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ffbd370cb847953a25954d9f63e14824a36113f8c72eecf6eccef5dc4b45d630"
-dependencies = [
- "crossbeam-channel",
-]
-
 [[package]]
 name = "thread_local"
 version = "1.1.9"
@@ -6776,17 +6358,6 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
 
-[[package]]
-name = "variadics_please"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "41b6d82be61465f97d42bd1d15bf20f3b0a3a0905018f38f9d6f6962055b0b5c"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.117",
-]
-
 [[package]]
 name = "variadics_please"
 version = "2.0.0"
@@ -7088,36 +6659,6 @@ dependencies = [
  "wezterm-dynamic",
 ]
 
-[[package]]
-name = "wgpu"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bb3feacc458f7bee8bc1737149b42b6c731aa461039a4264a67bb6681646b250"
-dependencies = [
- "arrayvec 0.7.6",
- "bitflags 2.11.1",
- "bytemuck",
- "cfg-if",
- "cfg_aliases",
- "document-features",
- "hashbrown 0.16.1",
- "js-sys",
- "log",
- "naga 29.0.3",
- "parking_lot",
- "portable-atomic",
- "profiling",
- "raw-window-handle",
- "smallvec",
- "static_assertions",
- "wasm-bindgen",
- "wasm-bindgen-futures",
- "web-sys",
- "wgpu-core 29.0.3",
- "wgpu-hal 29.0.3",
- "wgpu-types 29.0.3",
-]
-
 [[package]]
 name = "wgpu"
 version = "30.0.0"
@@ -7133,7 +6674,7 @@ dependencies = [
  "hashbrown 0.17.1",
  "js-sys",
  "log",
- "naga 30.0.0",
+ "naga",
  "parking_lot",
  "portable-atomic",
  "profiling",
@@ -7143,42 +6684,9 @@ dependencies = [
  "wasm-bindgen",
  "wasm-bindgen-futures",
  "web-sys",
- "wgpu-core 30.0.0",
- "wgpu-hal 30.0.0",
- "wgpu-types 30.0.1",
-]
-
-[[package]]
-name = "wgpu-core"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "02da3ad1b568337f25513b317870960ef87073ea0945502e44b864b67a8c77b7"
-dependencies = [
- "arrayvec 0.7.6",
- "bit-set 0.9.1",
- "bit-vec 0.9.1",
- "bitflags 2.11.1",
- "bytemuck",
- "cfg_aliases",
- "document-features",
- "hashbrown 0.16.1",
- "indexmap",
- "log",
- "naga 29.0.3",
- "once_cell",
- "parking_lot",
- "portable-atomic",
- "profiling",
- "raw-window-handle",
- "rustc-hash 1.1.0",
- "smallvec",
- "thiserror 2.0.18",
- "wgpu-core-deps-apple 29.0.3",
- "wgpu-core-deps-emscripten 29.0.3",
- "wgpu-core-deps-windows-linux-android 29.0.3",
- "wgpu-hal 29.0.3",
- "wgpu-naga-bridge 29.0.3",
- "wgpu-types 29.0.3",
+ "wgpu-core",
+ "wgpu-hal",
+ "wgpu-types",
 ]
 
 [[package]]
@@ -7197,7 +6705,7 @@ dependencies = [
  "hashbrown 0.17.1",
  "indexmap",
  "log",
- "naga 30.0.0",
+ "naga",
  "naga-types",
  "once_cell",
  "parking_lot",
@@ -7207,21 +6715,12 @@ dependencies = [
  "rustc-hash 1.1.0",
  "smallvec",
  "thiserror 2.0.18",
- "wgpu-core-deps-apple 30.0.1",
- "wgpu-core-deps-emscripten 30.0.1",
- "wgpu-core-deps-windows-linux-android 30.0.1",
- "wgpu-hal 30.0.0",
- "wgpu-naga-bridge 30.0.1",
- "wgpu-types 30.0.1",
-]
-
-[[package]]
-name = "wgpu-core-deps-apple"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "62e51b5447e144b3dbba4feb01f80f4fa21696fa0cd99afb2c3df1affd6fdb28"
-dependencies = [
- "wgpu-hal 29.0.3",
+ "wgpu-core-deps-apple",
+ "wgpu-core-deps-emscripten",
+ "wgpu-core-deps-windows-linux-android",
+ "wgpu-hal",
+ "wgpu-naga-bridge",
+ "wgpu-types",
 ]
 
 [[package]]
@@ -7230,16 +6729,7 @@ version = "30.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "061f3d319a40d39d00b1ecc2c33b89fe21d4e6fe01859df3500a3a8ecccd6b68"
 dependencies = [
- "wgpu-hal 30.0.0",
-]
-
-[[package]]
-name = "wgpu-core-deps-emscripten"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3487cd6293a963bc5c0c0396f6a2192043c50003c07f4efdccbad3d90ec9d819"
-dependencies = [
- "wgpu-hal 29.0.3",
+ "wgpu-hal",
 ]
 
 [[package]]
@@ -7248,16 +6738,7 @@ version = "30.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d98b86cf4abf524a902dd35f18ca6a3f08fc2ae9847c8f10b48e30491b1f0b86"
 dependencies = [
- "wgpu-hal 30.0.0",
-]
-
-[[package]]
-name = "wgpu-core-deps-windows-linux-android"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1bfb01076d0aa08b0ba9bd741e178b5cc440f5abe99d9581323a4c8b5d1a1916"
-dependencies = [
- "wgpu-hal 29.0.3",
+ "wgpu-hal",
 ]
 
 [[package]]
@@ -7266,61 +6747,7 @@ version = "30.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7586165fd5f6d881cb9ce4bb71f40d6caab2c0f1837e3fc1d9788a197fb6004f"
 dependencies = [
- "wgpu-hal 30.0.0",
-]
-
-[[package]]
-name = "wgpu-hal"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "31f8e1a9e7a8512f276f7c62e018c7fa8d60954303fed2e5750114332049193f"
-dependencies = [
- "android_system_properties",
- "arrayvec 0.7.6",
- "ash",
- "bit-set 0.9.1",
- "bitflags 2.11.1",
- "block2",
- "bytemuck",
- "cfg-if",
- "cfg_aliases",
- "glow",
- "glutin_wgl_sys",
- "gpu-allocator",
- "gpu-descriptor",
- "hashbrown 0.16.1",
- "js-sys",
- "khronos-egl",
- "libc",
- "libloading 0.8.9",
- "log",
- "naga 29.0.3",
- "ndk-sys",
- "objc2",
- "objc2-core-foundation",
- "objc2-foundation",
- "objc2-metal",
- "objc2-quartz-core",
- "once_cell",
- "ordered-float 5.3.0",
- "parking_lot",
- "portable-atomic",
- "portable-atomic-util",
- "profiling",
- "range-alloc",
- "raw-window-handle",
- "raw-window-metal",
- "renderdoc-sys",
- "smallvec",
- "thiserror 2.0.18",
- "wasm-bindgen",
- "wayland-sys",
- "web-sys",
- "wgpu-naga-bridge 29.0.3",
- "wgpu-types 29.0.3",
- "windows",
- "windows-core",
- "windows-result",
+ "wgpu-hal",
 ]
 
 [[package]]
@@ -7347,7 +6774,7 @@ dependencies = [
  "libc",
  "libloading 0.8.9",
  "log",
- "naga 30.0.0",
+ "naga",
  "naga-types",
  "ndk-sys",
  "objc2",
@@ -7372,45 +6799,21 @@ dependencies = [
  "wasm-bindgen",
  "wayland-sys",
  "web-sys",
- "wgpu-naga-bridge 30.0.1",
- "wgpu-types 30.0.1",
+ "wgpu-naga-bridge",
+ "wgpu-types",
  "windows",
  "windows-core",
  "windows-result",
 ]
 
-[[package]]
-name = "wgpu-naga-bridge"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "59c654c483f058800972c3645e95388a7eca31bf9fe1933bc20e036588a0be02"
-dependencies = [
- "naga 29.0.3",
- "wgpu-types 29.0.3",
-]
-
 [[package]]
 name = "wgpu-naga-bridge"
 version = "30.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d2f62e73117bb7a62bfd9c5a5841438a823f6566c6442a808ee269d2d055c081"
 dependencies = [
- "naga 30.0.0",
- "wgpu-types 30.0.1",
-]
-
-[[package]]
-name = "wgpu-types"
-version = "29.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a9bcc31518a0e9735aefebedb5f7a9ef3ed1c42549c9f4c882fa9060ceaac639"
-dependencies = [
- "bitflags 2.11.1",
- "bytemuck",
- "js-sys",
- "log",
- "raw-window-handle",
- "web-sys",
+ "naga",
+ "wgpu-types",
 ]
 
 [[package]]
diff --git a/engine/Cargo.toml b/engine/Cargo.toml
index 56d8f70..633fb12 100644
--- a/engine/Cargo.toml
+++ b/engine/Cargo.toml
@@ -6,7 +6,6 @@ edition = "2024"
 [dependencies]
 chess = "3.2.0"
 burn = { version = "0.22.0-pre.2", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
-burn-ndarray = "0.21.0"
 rand_distr = "0.6.0"
 rand = "0.10.1"
 ctrlc = "3.5.2"
diff --git a/engine/src/main.rs b/engine/src/main.rs
index 50a8a21..0000800 100644
--- a/engine/src/main.rs
+++ b/engine/src/main.rs
@@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig};
 fn main() {
     let device = Device::cuda(DeviceIndex::default());
 
-    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
+    let mcts_config = MctsConfig::new(100, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 
@@ -37,7 +37,7 @@ fn main() {
         seed: None,
     };
 
-    train(training_config, &device);
+    train(training_config, device);
 
     // let mut board_state: Option<BoardState> = Some(BoardState::new(
     //     Board::default(),
diff --git a/engine/src/train.rs b/engine/src/train.rs
index 625ec5a..d41c354 100644
--- a/engine/src/train.rs
+++ b/engine/src/train.rs
@@ -6,7 +6,6 @@ use crate::net::model::{
 use burn::data::dataloader::batcher::Batcher;
 use burn::module::{AutodiffModule, Module};
 use burn::optim::{AdamConfig, GradientsParams, Optimizer};
-// use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
 use burn::prelude::Device;
 use burn::store::ModuleRecord;
 use chess::{ChessMove, Color};
@@ -38,20 +37,32 @@ pub struct TrainingConfig {
     pub seed: Option<u64>,
 }
 
-pub fn train(training_config: TrainingConfig, device: &Device) {
+pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
+    let device = device.into();
+    if training_config.seed.is_some() {
+        device.seed(training_config.seed.unwrap());
+    }
+    let autodiff_device = device.clone().autodiff();
+
     let model_path = format!("artifacts/{}", training_config.model_name.as_str());
     let metadata_path = format!("{}.json", model_path);
     println!("Creating model...");
     let mut model: ChessModel = ChessModelConfig::init(
         training_config.num_blocks,
         training_config.hidden_channels,
-        device,
+        &autodiff_device,
     );
+    // Initialize optimizer once so state (moments) persist across steps
+    let mut optim = training_config.optimizer.init();
+
     if training_config.load_model {
         println!("Loading model {}...", model_path);
         // Load model in full precision from MessagePack file
         let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
-        model = model.load_record(record)
+        model = model.load_record(record);
+        optim = optim
+            .load(format!("{model_path}_optim"))
+            .expect("Failed to load optimizer");
     }
 
     let train = Arc::new(AtomicBool::new(true));
@@ -92,9 +103,6 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
     println!("using seed: {}", seed);
     let mut rng = SmallRng::seed_from_u64(seed);
 
-    // Initialize optimizer once so state (moments) persist across steps
-    let mut optim = training_config.optimizer.init();
-
     println!("Starting training...");
     while train.load(Ordering::Relaxed) {
         println!("Iteration: {}", iter);
@@ -116,11 +124,7 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
                 // println!("{}", before.elapsed().as_millis());
                 episode_buffer.push(results);
 
-                let temp = if ply_count < 30 {
-                    1.0
-                } else {
-                    0.0
-                };
+                let temp = if ply_count < 30 { 1.0 } else { 0.0 };
 
                 let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
 
@@ -201,8 +205,18 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
         let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
 
         let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
-        let p_loss_val = output.policy_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
-        let v_loss_val = output.value_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+        let p_loss_val = output
+            .policy_loss
+            .clone()
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
+        let v_loss_val = output
+            .value_loss
+            .clone()
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
 
         println!(
             "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
@@ -254,6 +268,10 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
         .into_record()
         .save(&model_path)
         .expect("Should be able to save the model");
+
+    optim
+        .save(format!("{model_path}_optim"))
+        .expect("Should be able to save the optim");
     println!("Model saved in {:?}, exiting training.", model_path);
     return;
 }