diff --git a/crates/burn-core/src/optim/adagrad.rs b/crates/burn-core/src/optim/adagrad.rs
index e64c75db6e..2df43619a6 100644
--- a/crates/burn-core/src/optim/adagrad.rs
+++ b/crates/burn-core/src/optim/adagrad.rs
@@ -27,9 +27,9 @@ pub struct AdaGradConfig {
 
 /// AdaGrad optimizer
 #[derive(Clone)]
-pub struct AdaGrad<B: Backend> {
+pub struct AdaGrad {
     lr_decay: LrDecay,
-    weight_decay: Option<WeightDecay<B>>,
+    weight_decay: Option<WeightDecay>,
 }
 
 /// AdaGrad state.
@@ -38,7 +38,7 @@ pub struct AdaGradState<B: Backend, const D: usize> {
     lr_decay: LrDecayState<B, D>,
 }
 
-impl<B: Backend> SimpleOptimizer<B> for AdaGrad<B> {
+impl<B: Backend> SimpleOptimizer<B> for AdaGrad {
     type State<const D: usize> = AdaGradState<B, D>;
 
     fn step<const D: usize>(
@@ -79,7 +79,7 @@ impl AdaGradConfig {
     /// Returns an optimizer that can be used to optimize a module.
     pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
         &self,
-    ) -> OptimizerAdaptor<AdaGrad<B::InnerBackend>, M, B> {
+    ) -> OptimizerAdaptor<AdaGrad, M, B> {
         let optim = AdaGrad {
             lr_decay: LrDecay {
                 lr_decay: self.lr_decay,
@@ -157,7 +157,7 @@ mod tests {
     use crate::optim::{GradientsParams, Optimizer};
     use crate::record::{BinFileRecorder, FullPrecisionSettings, Recorder};
     use crate::tensor::{Distribution, Tensor, TensorData};
-    use crate::{nn, nn::Linear, TestAutodiffBackend, TestBackend};
+    use crate::{nn, nn::Linear, TestAutodiffBackend};
 
     const LEARNING_RATE: LearningRate = 0.01;
 
@@ -274,8 +274,7 @@ mod tests {
     }
 
     fn create_adagrad(
-    ) -> OptimizerAdaptor<AdaGrad<TestBackend>, Linear<TestAutodiffBackend>, TestAutodiffBackend>
-    {
+    ) -> OptimizerAdaptor<AdaGrad, Linear<TestAutodiffBackend>, TestAutodiffBackend> {
         let config = AdaGradConfig::new();
         AdaGrad {
             lr_decay: LrDecay {
diff --git a/crates/burn-core/src/optim/adam.rs b/crates/burn-core/src/optim/adam.rs
index d7b4fe59b4..c2a5e3f679 100644
--- a/crates/burn-core/src/optim/adam.rs
+++ b/crates/burn-core/src/optim/adam.rs
@@ -32,18 +32,19 @@ pub struct AdamConfig {
 
 /// Adam optimizer as described in the paper [Adam: A Method for Stochastic Optimization](https://arxiv.org/pdf/1412.6980.pdf).
 #[derive(Clone)]
-pub struct Adam<B: Backend> {
+pub struct Adam {
     momentum: AdaptiveMomentum,
-    weight_decay: Option<WeightDecay<B>>,
+    weight_decay: Option<WeightDecay>,
 }
 
 /// Adam state.
 #[derive(Record, Clone, new)]
 pub struct AdamState<B: Backend, const D: usize> {
-    momentum: AdaptiveMomentumState<B, D>,
+    /// The current adaptive momentum.
+    pub momentum: AdaptiveMomentumState<B, D>,
 }
 
-impl<B: Backend> SimpleOptimizer<B> for Adam<B> {
+impl<B: Backend> SimpleOptimizer<B> for Adam {
     type State<const D: usize> = AdamState<B, D>;
 
     fn step<const D: usize>(
@@ -83,9 +84,7 @@ impl AdamConfig {
     /// # Returns
     ///
     /// Returns an optimizer that can be used to optimize a module.
-    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
-        &self,
-    ) -> OptimizerAdaptor<Adam<B::InnerBackend>, M, B> {
+    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(&self) -> OptimizerAdaptor<Adam, M, B> {
         let optim = Adam {
             momentum: AdaptiveMomentum {
                 beta_1: self.beta_1,
@@ -106,9 +105,12 @@ impl AdamConfig {
 /// Adaptive momentum state.
 #[derive(Record, new, Clone)]
 pub struct AdaptiveMomentumState<B: Backend, const D: usize> {
-    time: usize,
-    moment_1: Tensor<B, D>,
-    moment_2: Tensor<B, D>,
+    /// The number of iterations aggregated.
+    pub time: usize,
+    /// The first order momentum.
+    pub moment_1: Tensor<B, D>,
+    /// The second order momentum.
+    pub moment_2: Tensor<B, D>,
 }
 
 #[derive(Clone)]
@@ -190,7 +192,7 @@ mod tests {
     use crate::optim::{GradientsParams, Optimizer};
     use crate::record::{BinFileRecorder, FullPrecisionSettings, Recorder};
     use crate::tensor::{Distribution, Tensor, TensorData};
-    use crate::{nn, TestAutodiffBackend, TestBackend};
+    use crate::{nn, TestAutodiffBackend};
 
     const LEARNING_RATE: LearningRate = 0.01;
 
@@ -350,8 +352,7 @@ mod tests {
             .load_record(record)
     }
 
-    fn create_adam(
-    ) -> OptimizerAdaptor<Adam<TestBackend>, nn::Linear<TestAutodiffBackend>, TestAutodiffBackend>
+    fn create_adam() -> OptimizerAdaptor<Adam, nn::Linear<TestAutodiffBackend>, TestAutodiffBackend>
     {
         let config = AdamConfig::new();
         Adam {
diff --git a/crates/burn-core/src/optim/adamw.rs b/crates/burn-core/src/optim/adamw.rs
index 666ec18f13..32f4898dc6 100644
--- a/crates/burn-core/src/optim/adamw.rs
+++ b/crates/burn-core/src/optim/adamw.rs
@@ -1,14 +1,12 @@
+use super::{AdaptiveMomentumState, SimpleOptimizer};
+use crate::config::Config;
+use crate::optim::adaptor::OptimizerAdaptor;
+use crate::tensor::{backend::AutodiffBackend, Tensor};
 use crate::{
     self as burn, grad_clipping::GradientClippingConfig, module::AutodiffModule, record::Record,
     LearningRate,
 };
-use std::marker::PhantomData;
-
-use super::SimpleOptimizer;
-use crate::config::Config;
-use crate::optim::adaptor::OptimizerAdaptor;
-use crate::tensor::{backend::AutodiffBackend, Tensor};
-use burn_tensor::{backend::Backend, ops::Device, ElementConversion};
+use burn_tensor::{backend::Backend, ops::Device};
 
 /// AdamW configuration.
 #[derive(Config)]
@@ -31,19 +29,19 @@ pub struct AdamWConfig {
 
 /// AdamW optimizer as described in the paper [Decoupled Weight Decay Regularization, Loshchilov and Hutter, 2019](https://arxiv.org/abs/1711.05101).
 #[derive(Clone)]
-pub struct AdamW<B: Backend> {
+pub struct AdamW {
     momentum: AdaptiveMomentumW,
     weight_decay: f32,
-    _phantom: PhantomData<B>,
 }
 
 /// AdamW state.
 #[derive(Record, Clone, new)]
 pub struct AdamWState<B: Backend, const D: usize> {
-    momentum: AdaptiveMomentumWState<B, D>,
+    /// Th current adaptive momentum state.
+    pub momentum: AdaptiveMomentumState<B, D>,
 }
 
-impl<B: Backend> SimpleOptimizer<B> for AdamW<B> {
+impl<B: Backend> SimpleOptimizer<B> for AdamW {
     type State<const D: usize> = AdamWState<B, D>;
 
     /// A single optimization step for any tensor that represents the parameters of a model.
@@ -81,9 +79,7 @@ impl AdamWConfig {
     /// # Returns
     ///
     /// Returns an optimizer that can be used to optimize a module.
-    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
-        &self,
-    ) -> OptimizerAdaptor<AdamW<B::InnerBackend>, M, B> {
+    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(&self) -> OptimizerAdaptor<AdamW, M, B> {
         let optim = AdamW {
             momentum: AdaptiveMomentumW {
                 beta_1: self.beta_1,
@@ -91,7 +87,6 @@ impl AdamWConfig {
                 epsilon: self.epsilon,
             },
             weight_decay: self.weight_decay,
-            _phantom: Default::default(),
         };
 
         let mut optim = OptimizerAdaptor::from(optim);
@@ -102,14 +97,6 @@ impl AdamWConfig {
     }
 }
 
-/// Adaptive momentum state.
-#[derive(Record, new, Clone)]
-pub struct AdaptiveMomentumWState<B: Backend, const D: usize> {
-    time: usize,
-    moment_1: Tensor<B, D>,
-    moment_2: Tensor<B, D>,
-}
-
 #[derive(Clone)]
 struct AdaptiveMomentumW {
     beta_1: f32,
@@ -121,8 +108,8 @@ impl AdaptiveMomentumW {
     pub fn transform<B: Backend, const D: usize>(
         &self,
         grad: Tensor<B, D>,
-        state: Option<AdaptiveMomentumWState<B, D>>,
-    ) -> (Tensor<B, D>, AdaptiveMomentumWState<B, D>) {
+        state: Option<AdaptiveMomentumState<B, D>>,
+    ) -> (Tensor<B, D>, AdaptiveMomentumState<B, D>) {
         let state = if let Some(mut state) = state {
             // Update first moment estimate.
             let factor = 1.0 - self.beta_1;
@@ -151,10 +138,10 @@ impl AdaptiveMomentumW {
             let factor = 1.0 - self.beta_2;
             let moment_2 = grad.powf_scalar(2.0).mul_scalar(factor);
 
-            AdaptiveMomentumWState::new(1, moment_1, moment_2)
+            AdaptiveMomentumState::new(1, moment_1, moment_2)
         };
 
-        let time: i32 = (state.time as i32).elem();
+        let time: i32 = state.time as i32;
 
         // Compute bias-corrected first and second moment estimates.
         let moment_1_corrected = state
@@ -173,28 +160,11 @@ impl AdaptiveMomentumW {
 
         (
             update_delta,
-            AdaptiveMomentumWState::new(state.time, state.moment_1, state.moment_2),
+            AdaptiveMomentumState::new(state.time, state.moment_1, state.moment_2),
         )
     }
 }
 
-impl<B: Backend, const D: usize> AdaptiveMomentumWState<B, D> {
-    /// Move state to device.
-    ///
-    /// # Arguments
-    ///
-    /// * `device` - Device to move state to.
-    ///
-    /// # Returns
-    ///
-    /// Returns state moved to device.
-    pub fn to_device(mut self, device: &B::Device) -> Self {
-        self.moment_1 = self.moment_1.to_device(device);
-        self.moment_2 = self.moment_2.to_device(device);
-        self
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -202,7 +172,7 @@ mod tests {
     use crate::optim::{GradientsParams, Optimizer};
     use crate::record::{BinFileRecorder, FullPrecisionSettings, Recorder};
     use crate::tensor::{Distribution, Tensor, TensorData};
-    use crate::{nn, TestAutodiffBackend, TestBackend};
+    use crate::{nn, TestAutodiffBackend};
     use tempfile::TempDir;
 
     const LEARNING_RATE: LearningRate = 0.01;
@@ -366,8 +336,7 @@ mod tests {
     }
 
     fn create_adamw(
-    ) -> OptimizerAdaptor<AdamW<TestBackend>, nn::Linear<TestAutodiffBackend>, TestAutodiffBackend>
-    {
+    ) -> OptimizerAdaptor<AdamW, nn::Linear<TestAutodiffBackend>, TestAutodiffBackend> {
         let config = AdamWConfig::new();
         AdamW {
             momentum: AdaptiveMomentumW {
@@ -376,7 +345,6 @@ mod tests {
                 epsilon: config.epsilon,
             },
             weight_decay: config.weight_decay,
-            _phantom: Default::default(),
         }
         .into()
     }
diff --git a/crates/burn-core/src/optim/decay.rs b/crates/burn-core/src/optim/decay.rs
index 9979f57c29..27f258c0d8 100644
--- a/crates/burn-core/src/optim/decay.rs
+++ b/crates/burn-core/src/optim/decay.rs
@@ -4,13 +4,13 @@ use crate as burn;
 use crate::record::Record;
 
 use crate::config::Config;
-use crate::tensor::{ElementConversion, Tensor};
+use crate::tensor::Tensor;
 
 /// Configuration to create [weight decay](WeightDecay).
 #[derive(Config)]
 pub struct WeightDecayConfig {
     /// L2 penalty.
-    pub penalty: f64,
+    pub penalty: f32,
 }
 
 /// State of [weight decay](WeightDecay).
@@ -21,15 +21,15 @@ pub struct WeightDecayState<B: Backend, const D: usize> {
 
 /// Weight decay implementation that transforms gradients.
 #[derive(Clone)]
-pub struct WeightDecay<B: Backend> {
-    penalty: B::FloatElem,
+pub struct WeightDecay {
+    penalty: f32,
 }
 
-impl<B: Backend> WeightDecay<B> {
+impl WeightDecay {
     /// Creates a new [weight decay](WeightDecay) from a [config](WeightDecayConfig).
     pub fn new(config: &WeightDecayConfig) -> Self {
         Self {
-            penalty: config.penalty.elem(),
+            penalty: config.penalty,
         }
     }
 
@@ -43,7 +43,7 @@ impl<B: Backend> WeightDecay<B> {
     /// # Returns
     ///
     /// * `grad` - Transformed gradient.
-    pub fn transform<const D: usize>(
+    pub fn transform<B: Backend, const D: usize>(
         &self,
         grad: Tensor<B, D>,
         tensor: Tensor<B, D>,
diff --git a/crates/burn-core/src/optim/rmsprop.rs b/crates/burn-core/src/optim/rmsprop.rs
index 63ccbe542d..bff9859df5 100644
--- a/crates/burn-core/src/optim/rmsprop.rs
+++ b/crates/burn-core/src/optim/rmsprop.rs
@@ -41,7 +41,7 @@ impl RmsPropConfig {
     /// Returns an optimizer that can be used to optimize a module.
     pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
         &self,
-    ) -> OptimizerAdaptor<RmsProp<B::InnerBackend>, M, B> {
+    ) -> OptimizerAdaptor<RmsProp, M, B> {
         let weight_decay = self.weight_decay.as_ref().map(WeightDecay::new);
 
         let mut optim = OptimizerAdaptor::from(RmsProp {
@@ -65,16 +65,16 @@ impl RmsPropConfig {
 /// Optimizer that implements stochastic gradient descent with momentum.
 /// The optimizer can be configured with [RmsPropConfig](RmsPropConfig).
 #[derive(Clone)]
-pub struct RmsProp<B: Backend> {
+pub struct RmsProp {
     alpha: f32,
     // epsilon: f32,
     centered: bool,
     // momentum: Option<Momentum<B>>,
     momentum: RmsPropMomentum,
-    weight_decay: Option<WeightDecay<B>>,
+    weight_decay: Option<WeightDecay>,
 }
 
-impl<B: Backend> SimpleOptimizer<B> for RmsProp<B> {
+impl<B: Backend> SimpleOptimizer<B> for RmsProp {
     type State<const D: usize> = RmsPropState<B, D>;
 
     fn step<const D: usize>(
@@ -136,15 +136,19 @@ impl<B: Backend> SimpleOptimizer<B> for RmsProp<B> {
 /// State of [RmsProp](RmsProp)
 #[derive(Record, Clone, new)]
 pub struct RmsPropState<B: Backend, const D: usize> {
-    square_avg: SquareAvgState<B, D>,
-    centered: CenteredState<B, D>,
-    momentum: Option<RmsPropMomentumState<B, D>>,
+    /// Current squared average state.
+    pub square_avg: SquareAvgState<B, D>,
+    /// Current centered state
+    pub centered: CenteredState<B, D>,
+    /// Current gradient momentum, if any.
+    pub momentum: Option<RmsPropMomentumState<B, D>>,
 }
 
 /// [SquareAvgState](SquareAvgState) is to store and pass optimizer step params.
 #[derive(Record, Clone, new)]
 pub struct SquareAvgState<B: Backend, const D: usize> {
-    square_avg: Tensor<B, D>,
+    /// Current squared average.
+    pub square_avg: Tensor<B, D>,
 }
 
 impl<B: Backend, const D: usize> SquareAvgState<B, D> {
@@ -183,8 +187,10 @@ impl<B: Backend, const D: usize> SquareAvgState<B, D> {
 /// [CenteredState](CenteredState) is to store and pass optimizer step params.
 #[derive(Record, Clone, new)]
 pub struct CenteredState<B: Backend, const D: usize> {
-    grad_avg: Option<Tensor<B, D>>,
-    avg: Tensor<B, D>,
+    /// The averaged gradient to calculate the centered gradient, if available.
+    pub grad_avg: Option<Tensor<B, D>>,
+    /// The current average value.
+    pub avg: Tensor<B, D>,
 }
 
 impl<B: Backend, const D: usize> CenteredState<B, D> {
@@ -316,7 +322,7 @@ mod tests {
     use crate::optim::{GradientsParams, Optimizer};
     use crate::record::{BinFileRecorder, FullPrecisionSettings, Recorder};
     use crate::tensor::{Distribution, Tensor, TensorData};
-    use crate::{nn, TestAutodiffBackend, TestBackend};
+    use crate::{nn, TestAutodiffBackend};
     use tempfile::TempDir;
 
     const LEARNING_RATE: LearningRate = 0.01;
@@ -530,8 +536,7 @@ mod tests {
     }
 
     fn create_rmsprop(
-    ) -> OptimizerAdaptor<RmsProp<TestBackend>, nn::Linear<TestAutodiffBackend>, TestAutodiffBackend>
-    {
+    ) -> OptimizerAdaptor<RmsProp, nn::Linear<TestAutodiffBackend>, TestAutodiffBackend> {
         RmsPropConfig {
             alpha: 0.99,
             epsilon: 1e-9,
diff --git a/crates/burn-core/src/optim/sgd.rs b/crates/burn-core/src/optim/sgd.rs
index 325325cbfa..9b5d66c4fa 100644
--- a/crates/burn-core/src/optim/sgd.rs
+++ b/crates/burn-core/src/optim/sgd.rs
@@ -28,13 +28,14 @@ pub struct SgdConfig {
 #[derive(Clone)]
 pub struct Sgd<B: Backend> {
     momentum: Option<Momentum<B>>,
-    weight_decay: Option<WeightDecay<B>>,
+    weight_decay: Option<WeightDecay>,
 }
 
 /// State of [Sgd](Sgd).
 #[derive(Record, Clone, new)]
 pub struct SgdState<B: Backend, const D: usize> {
-    momentum: Option<MomentumState<B, D>>,
+    /// The current state of the momentum (if any).
+    pub momentum: Option<MomentumState<B, D>>,
 }
 
 impl SgdConfig {