diff --git a/src/models/unet_2d_blocks.rs b/src/models/unet_2d_blocks.rs
index 22841c6..26b6a5f 100644
--- a/src/models/unet_2d_blocks.rs
+++ b/src/models/unet_2d_blocks.rs
@@ -567,10 +567,6 @@ pub struct CrossAttnDownBlock2D<B: Backend> {
 
 impl CrossAttnDownBlock2DConfig {
     pub fn init<B: Backend>(&self, device: &B::Device) -> CrossAttnDownBlock2D<B> {
-        let mut downblock = self.downblock;
-        downblock.in_channels = self.in_channels;
-        downblock.out_channels = self.out_channels;
-        downblock.temb_channels = self.temb_channels;
         let downblock = self.downblock.init(device);
 
         let attentions = (0..self.downblock.n_layers)
diff --git a/src/schedulers/ddim.rs b/src/schedulers/ddim.rs
index 3d78d5f..3140ecc 100644
--- a/src/schedulers/ddim.rs
+++ b/src/schedulers/ddim.rs
@@ -11,12 +11,12 @@
 use alloc::vec::Vec;
 use burn::tensor::{backend::Backend, Tensor};
 
-use super::{betas_for_alpha_bar, BetaSchedule, PredictionType};
-
 #[cfg(not(feature = "std"))]
 #[allow(unused_imports)]
 use num_traits::Float;
 
+use super::{betas_for_alpha_bar, BetaSchedule, PredictionType};
+
 /// The configuration for the DDIM scheduler.
 #[derive(Debug, Clone, Copy)]
 pub struct DDIMSchedulerConfig {
diff --git a/src/schedulers/dpmsolver_multistep.rs b/src/schedulers/dpmsolver_multistep.rs
index 5cb7066..d042288 100644
--- a/src/schedulers/dpmsolver_multistep.rs
+++ b/src/schedulers/dpmsolver_multistep.rs
@@ -11,12 +11,12 @@ use alloc::vec;
 use alloc::vec::Vec;
 use burn::tensor::{backend::Backend, Tensor};
 
-use super::{betas_for_alpha_bar, BetaSchedule, PredictionType};
-
 #[cfg(not(feature = "std"))]
 #[allow(unused_imports)]
 use num_traits::Float;
 
+use super::{betas_for_alpha_bar, BetaSchedule, PredictionType};
+
 /// The algorithm type for the solver.
 #[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
 pub enum DPMSolverAlgorithmType {
diff --git a/src/schedulers/euler_discrete.rs b/src/schedulers/euler_discrete.rs
index 6d8f792..3fb89c9 100644
--- a/src/schedulers/euler_discrete.rs
+++ b/src/schedulers/euler_discrete.rs
@@ -8,6 +8,10 @@ use alloc::vec;
 use alloc::vec::Vec;
 use burn::tensor::{backend::Backend, Distribution, Tensor};
 
+#[cfg(not(feature = "std"))]
+#[allow(unused_imports)]
+use num_traits::Float;
+
 use super::{BetaSchedule, PredictionType};
 
 /// Configuration for the Euler Discrete Scheduler.
diff --git a/src/schedulers/mod.rs b/src/schedulers/mod.rs
index 6ee8c80..75bc9dc 100644
--- a/src/schedulers/mod.rs
+++ b/src/schedulers/mod.rs
@@ -7,6 +7,10 @@ use alloc::vec::Vec;
 use burn::tensor::{backend::Backend, Tensor};
 use core::f64::consts::FRAC_PI_2;
 
+#[cfg(not(feature = "std"))]
+#[allow(unused_imports)]
+use num_traits::Float;
+
 pub mod ddim;
 pub mod ddpm;
 pub mod dpmsolver_multistep;