Fix for cubecl update

tracel-ai · Feb 13, 2025 · bbc92c6 · bbc92c6
1 parent 134aff7
commit bbc92c6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/crates/burn-jit/src/fusion/matmul/optimization.rs b/crates/burn-jit/src/fusion/matmul/optimization.rs
@@ -408,19 +408,19 @@ fn matmul_launch_kernel<'a, R: Runtime, EG: Numeric, S: MatmulSelector>(
         || TypeId::of::<EG>() == TypeId::of::<flex32>()
     {
         S::select_kernel::<FusedMatmulSpec<EG, half::f16, f32>, R>(
-            client, input, output, problem, plane_size,
+            client, input, output, problem, plane_size, false,
         )
     } else if TypeId::of::<EG>() == TypeId::of::<half::bf16>() {
         S::select_kernel::<FusedMatmulSpec<EG, half::bf16, f32>, R>(
-            client, input, output, problem, plane_size,
+            client, input, output, problem, plane_size, false,
         )
     } else if S::stage_tf32_supported() {
         S::select_kernel::<FusedMatmulSpec<EG, tf32, f32>, R>(
-            client, input, output, problem, plane_size,
+            client, input, output, problem, plane_size, false,
         )
     } else {
         S::select_kernel::<FusedMatmulSpec<EG, EG, f32>, R>(
-            client, input, output, problem, plane_size,
+            client, input, output, problem, plane_size, false,
         )
     }
 }
diff --git a/crates/burn-jit/src/kernel/conv/conv2d/gemm/homogeneous/base.rs b/crates/burn-jit/src/kernel/conv/conv2d/gemm/homogeneous/base.rs
@@ -207,6 +207,7 @@ where
             cube_dim,
             cube_count,
             advanced_config,
+            false,
         );
         let size = SMM::stage_shape(&smm_config);