diff --git a/Cargo.lock b/Cargo.lock
index 50b5f0af1..75794f592 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1397,7 +1397,7 @@ dependencies = [
 [[package]]
 name = "cubecl"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "cubecl-core",
  "cubecl-cuda",
@@ -1408,7 +1408,7 @@ dependencies = [
 [[package]]
 name = "cubecl-common"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "derive-new",
  "getrandom",
@@ -1423,7 +1423,7 @@ dependencies = [
 [[package]]
 name = "cubecl-core"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1433,13 +1433,14 @@ dependencies = [
  "half",
  "log",
  "num-traits",
+ "paste",
  "serde",
 ]
 
 [[package]]
 name = "cubecl-cuda"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1454,7 +1455,7 @@ dependencies = [
 [[package]]
 name = "cubecl-linalg"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "bytemuck",
  "cubecl-core",
@@ -1465,7 +1466,7 @@ dependencies = [
 [[package]]
 name = "cubecl-macros"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "cubecl-common",
  "darling",
@@ -1480,7 +1481,7 @@ dependencies = [
 [[package]]
 name = "cubecl-runtime"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "async-channel",
  "cfg_aliases 0.2.1",
@@ -1500,7 +1501,7 @@ dependencies = [
 [[package]]
 name = "cubecl-wgpu"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=7a86f9a86e376fedb09f096f2b548e501a130883#7a86f9a86e376fedb09f096f2b548e501a130883"
+source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
 dependencies = [
  "async-channel",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index 7de841f4c..b06407b59 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -151,8 +151,8 @@ tch = "0.15.0"
 portable-atomic-util = { version = "0.2.2", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "7a86f9a86e376fedb09f096f2b548e501a130883" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "7a86f9a86e376fedb09f096f2b548e501a130883" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl" }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common" }