diff --git a/Cargo.lock b/Cargo.lock
index 7f5ed430c..d09aed2c5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1395,7 +1395,7 @@ dependencies = [
 [[package]]
 name = "cubecl"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "cubecl-core",
  "cubecl-cuda",
@@ -1406,7 +1406,7 @@ dependencies = [
 [[package]]
 name = "cubecl-common"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "derive-new",
  "getrandom",
@@ -1421,7 +1421,7 @@ dependencies = [
 [[package]]
 name = "cubecl-core"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "bytemuck",
  "cubecl-macros",
@@ -1436,7 +1436,7 @@ dependencies = [
 [[package]]
 name = "cubecl-cuda"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1451,7 +1451,7 @@ dependencies = [
 [[package]]
 name = "cubecl-linalg"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "bytemuck",
  "cubecl-core",
@@ -1462,7 +1462,7 @@ dependencies = [
 [[package]]
 name = "cubecl-macros"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "derive-new",
  "proc-macro2",
@@ -1473,7 +1473,7 @@ dependencies = [
 [[package]]
 name = "cubecl-runtime"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "async-channel",
  "cfg_aliases 0.2.1",
@@ -1493,7 +1493,7 @@ dependencies = [
 [[package]]
 name = "cubecl-wgpu"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e#446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e"
+source = "git+https://github.com/tracel-ai/cubecl?rev=b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73#b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73"
 dependencies = [
  "async-channel",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index db8604e4f..27c376ecd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -151,8 +151,8 @@ systemstat = "0.2.3"
 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 = "446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "446dc854ef3f7f93fe4ab2df0f1ab779a7d6a98e" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "b0fa39bad9d5ecbd98a955e6f7f94c855b6e9a73" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl" }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common" }
diff --git a/backend-comparison/Cargo.toml b/backend-comparison/Cargo.toml
index 03d5d6f42..fd2a180b5 100644
--- a/backend-comparison/Cargo.toml
+++ b/backend-comparison/Cargo.toml
@@ -24,14 +24,13 @@ tch-cpu = ["burn/tch"]
 tch-gpu = ["burn/tch"]
 wgpu = ["burn/wgpu", "burn/autotune"]
 wgpu-fusion = ["wgpu", "burn/fusion"]
-# cuda-jit = ["burn-cuda"]
+cuda-jit = ["burn/cuda-jit"]
 
 [dependencies]
 arboard = { workspace = true }
 burn = { path = "../crates/burn", default-features = false }
 burn-common = { path = "../crates/burn-common", version = "0.15.0" }
 burn-wgpu = { path = "../crates/burn-wgpu", default-features = false, version = "0.15.0", optional = true }
-# burn-cuda = { path = "../crates/burn-cuda", version = "0.15.0", optional = true }
 clap = { workspace = true }
 colored = { workspace = true }
 derive-new = { workspace = true }
diff --git a/backend-comparison/src/lib.rs b/backend-comparison/src/lib.rs
index f6cefcd36..0f371fcc0 100644
--- a/backend-comparison/src/lib.rs
+++ b/backend-comparison/src/lib.rs
@@ -129,7 +129,7 @@ macro_rules! bench_on_backend {
 
         #[cfg(feature = "cuda-jit")]
         {
-            use burn_cuda::{Cuda, CudaDevice};
+            use burn::backend::cuda_jit::{Cuda, CudaDevice};
 
             bench::<Cuda>(&CudaDevice::default(), feature_name, url, token);
         }