diff --git a/Cargo.lock b/Cargo.lock
index b449b4dc5e..25f8f81667 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1406,7 +1406,7 @@ dependencies = [
 [[package]]
 name = "cubecl"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "cubecl-core",
  "cubecl-cuda",
@@ -1418,7 +1418,7 @@ dependencies = [
 [[package]]
 name = "cubecl-common"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "derive-new",
  "getrandom",
@@ -1433,7 +1433,7 @@ dependencies = [
 [[package]]
 name = "cubecl-core"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1450,7 +1450,7 @@ dependencies = [
 [[package]]
 name = "cubecl-cuda"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1465,7 +1465,7 @@ dependencies = [
 [[package]]
 name = "cubecl-linalg"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "bytemuck",
  "cubecl-core",
@@ -1476,7 +1476,7 @@ dependencies = [
 [[package]]
 name = "cubecl-macros"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "cubecl-common",
  "darling",
@@ -1491,7 +1491,7 @@ dependencies = [
 [[package]]
 name = "cubecl-runtime"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "async-channel",
  "cfg_aliases 0.2.1",
@@ -1511,7 +1511,7 @@ dependencies = [
 [[package]]
 name = "cubecl-wgpu"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=52e21ff91890c8b0ae4943816fdcaa7e3f4be058#52e21ff91890c8b0ae4943816fdcaa7e3f4be058"
+source = "git+https://github.com/tracel-ai/cubecl?rev=fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd#fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd"
 dependencies = [
  "async-channel",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index 0f44b8c8e7..5d74140d35 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -152,8 +152,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 = "52e21ff91890c8b0ae4943816fdcaa7e3f4be058" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "52e21ff91890c8b0ae4943816fdcaa7e3f4be058" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb47090f7a44952ae3e3b2b72f8c5a88d8af56fd" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl" }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common" }