matmul vec4

tracel-ai · Oct 24, 2023 · 5b3bfdb · 5b3bfdb
1 parent 0ab611b
commit 5b3bfdb
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/burn-wgpu/src/ops/float_ops.rs b/burn-wgpu/src/ops/float_ops.rs
@@ -1,4 +1,5 @@
 use super::{numeric, BoolTensor, Device, FloatElem, FloatTensor, FullPrecisionBackend, IntTensor};
+use crate::kernel::matmul::vec4_primitive::matmul_tiling_2d_vec4_primitive_default;
 use crate::kernel::prng::{random_bernoulli, random_normal, random_uniform};
 use crate::kernel::{
     self, unary_default, unary_inplace_default, unary_scalar_default, unary_scalar_inplace_default,
@@ -143,7 +144,7 @@ where
         lhs: FloatTensor<Self, D>,
         rhs: FloatTensor<Self, D>,
     ) -> FloatTensor<Self, D> {
-        kernel::matmul::contiguous::matmul_tiling_2d_default(lhs, rhs)
+        matmul_tiling_2d_vec4_primitive_default(lhs, rhs)
     }
 
     fn swap_dims<const D: usize>(

diff --git a/examples/mnist/src/training.rs b/examples/mnist/src/training.rs
@@ -3,7 +3,7 @@ use crate::model::Model;
 
 use burn::module::Module;
 use burn::optim::decay::WeightDecayConfig;
-use burn::optim::{AdamConfig, AdamWConfig};
+use burn::optim::AdamConfig;
 use burn::record::{CompactRecorder, NoStdTrainingRecorder};
 use burn::train::metric::store::{Aggregate, Direction, Split};
 use burn::train::metric::{CpuMemory, CpuTemperature, CpuUse};