diff --git a/crates/ml/src/hyperopt/adapters/mamba2.rs b/crates/ml/src/hyperopt/adapters/mamba2.rs
index cae31f9f2..833349fc5 100644
--- a/crates/ml/src/hyperopt/adapters/mamba2.rs
+++ b/crates/ml/src/hyperopt/adapters/mamba2.rs
@@ -617,13 +617,13 @@ impl Mamba2Trainer {
             // Normalize target to [0,1]
             let normalized_target = (target_price - target_min) / (target_max - target_min);
 
-            let input_tensor = Tensor::new(sequence.as_slice(), &Device::Cpu)?.reshape((
+            let input_tensor = Tensor::new(sequence.as_slice(), &self.device)?.reshape((
                 1,
                 seq_len,
                 self.d_model,
             ))?;
             let target_tensor =
-                Tensor::new(&[normalized_target], &Device::Cpu)?.reshape((1, 1, 1))?;
+                Tensor::new(&[normalized_target], &self.device)?.reshape((1, 1, 1))?;
 
             feature_sequences.push((input_tensor, target_tensor));
         }