diff --git a/crates/ml/src/dqn/curiosity.rs b/crates/ml/src/dqn/curiosity.rs
index 3e2c07c19..18a452e16 100644
--- a/crates/ml/src/dqn/curiosity.rs
+++ b/crates/ml/src/dqn/curiosity.rs
@@ -3,7 +3,7 @@
 //! Implements forward dynamics model that predicts next state from (state, action)
 //! and provides novelty-based intrinsic rewards via prediction error.
 
-use candle_core::{DType, Device, Tensor};
+use candle_core::{Device, Tensor};
 use candle_nn::{ops::leaky_relu, AdamW, Linear, Module, Optimizer, ParamsAdamW, VarBuilder, VarMap};
 
 use super::action_space::{FactoredAction, ExposureLevel};
@@ -70,12 +70,12 @@ impl ForwardDynamicsModel {
         // Extract first 32 features from state (state embedding)
         let state_embedding = state.narrow(1, 0, 32)
             .map_err(|e| MLError::ModelError(format!("Failed to narrow state: {}", e)))?
-            .to_dtype(DType::F32)
-            .map_err(|e| MLError::ModelError(format!("Failed to convert state to F32: {}", e)))?;
+            .to_dtype(training_dtype(&self.device))
+            .map_err(|e| MLError::ModelError(format!("Failed to convert state dtype: {}", e)))?;
 
         // One-hot encode action (convert FactoredAction to simplified action index)
         let batch_size = state.dims()[0];
-        let mut action_onehot = Tensor::zeros((batch_size, 3), DType::F32, &self.device)
+        let mut action_onehot = Tensor::zeros((batch_size, 3), training_dtype(&self.device), &self.device)
             .map_err(|e| MLError::ModelError(format!("Failed to create action tensor: {}", e)))?;
 
         // Convert FactoredAction to simplified action index: 0=BUY, 1=SELL, 2=HOLD
@@ -85,7 +85,7 @@ impl ForwardDynamicsModel {
             ExposureLevel::Flat => 2_i64, // HOLD
         };
         for batch_idx in 0..batch_size {
-            action_onehot = action_onehot.slice_assign(&[batch_idx..batch_idx+1, action_idx as usize..action_idx as usize+1], &Tensor::ones((1, 1), DType::F32, &self.device)?)
+            action_onehot = action_onehot.slice_assign(&[batch_idx..batch_idx+1, action_idx as usize..action_idx as usize+1], &Tensor::ones((1, 1), training_dtype(&self.device), &self.device)?)
                 .map_err(|e| MLError::ModelError(format!("Failed to set action one-hot: {}", e)))?;
         }
 
@@ -201,8 +201,8 @@ impl CuriosityModule {
         // Extract next state embedding (first 32 features)
         let next_state_embedding = next_state.narrow(1, 0, 32)
             .map_err(|e| MLError::ModelError(format!("Failed to narrow next_state: {}", e)))?
-            .to_dtype(DType::F32)
-            .map_err(|e| MLError::ModelError(format!("Failed to convert next_state to F32: {}", e)))?;
+            .to_dtype(training_dtype(&self.forward_model.device))
+            .map_err(|e| MLError::ModelError(format!("Failed to convert next_state dtype: {}", e)))?;
 
         // Predict next state
         let predicted_next_state = self.forward_model.predict(state, action)?;
@@ -235,7 +235,7 @@ impl CuriosityModule {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use candle_core::Device;
+    use candle_core::{DType, Device};
     use super::super::action_space::{FactoredAction, ExposureLevel, OrderType, Urgency};
 
     // Helper to create a test BUY action