From 129c2f19b83e620b2270bb47c9a308ae4e9c0e34 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 2 Mar 2026 18:03:10 +0100 Subject: [PATCH] fix(ml): mamba2 hyperopt tensors created on GPU instead of CPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace &Device::Cpu with &self.device for input and target tensor creation in Mamba2 hyperopt adapter. Avoids unnecessary CPU→GPU transfer during hyperparameter optimization. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/hyperopt/adapters/mamba2.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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)); }