diff --git a/ml/src/mamba/selective_state.rs b/ml/src/mamba/selective_state.rs
index 8f17a59df..b88e80485 100644
--- a/ml/src/mamba/selective_state.rs
+++ b/ml/src/mamba/selective_state.rs
@@ -16,7 +16,10 @@ use std::collections::{BTreeMap, HashMap, VecDeque};
 use std::mem::size_of;
 use std::sync::atomic::{AtomicU64, Ordering};
 
+#[cfg(not(test))]
 use candle_core::Tensor;
+#[cfg(test)]
+use candle_core::{Device, Tensor};
 use nalgebra::DVector;
 use tracing::{debug, instrument};
 
diff --git a/ml/src/tft/quantized_vsn.rs b/ml/src/tft/quantized_vsn.rs
index 4f731110f..fd2a08cb3 100644
--- a/ml/src/tft/quantized_vsn.rs
+++ b/ml/src/tft/quantized_vsn.rs
@@ -12,9 +12,14 @@ use candle_nn::{VarBuilder, VarMap};
 use tracing::{debug, info};
 
 use super::variable_selection::VariableSelectionNetwork;
+#[cfg(not(test))]
 use crate::memory_optimization::quantization::{
     QuantizationConfig, Quantizer, QuantizedTensor,
 };
+#[cfg(test)]
+use crate::memory_optimization::quantization::{
+    QuantizationConfig, QuantizationType, Quantizer, QuantizedTensor,
+};
 use crate::MLError;
 
 /// Quantized Variable Selection Network