diff --git a/crates/ml/src/ensemble/adapters/dqn.rs b/crates/ml/src/ensemble/adapters/dqn.rs
index 21a6bfaff..062de4b08 100644
--- a/crates/ml/src/ensemble/adapters/dqn.rs
+++ b/crates/ml/src/ensemble/adapters/dqn.rs
@@ -303,7 +303,7 @@ mod tests {
         let _guard = CUDA_LOCK.lock().unwrap_or_else(|p| p.into_inner());
         let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap();
         let fv = FeatureVector {
-            values: vec![0.1; 56],
+            values: vec![0.1; ml_core::state_layout::STATE_DIM],
             timestamp: 1700000000_000_000,
         };
         let pred = adapter.predict(&fv).unwrap();
@@ -328,7 +328,7 @@ mod tests {
         let _guard = CUDA_LOCK.lock().unwrap_or_else(|p| p.into_inner());
         let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap();
         let fv = FeatureVector {
-            values: vec![0.1; 56],
+            values: vec![0.1; ml_core::state_layout::STATE_DIM],
             timestamp: 1700000000_000_000,
         };
         let pred1 = adapter.predict(&fv).unwrap();
@@ -361,7 +361,7 @@ mod tests {
         }
 
         let fv = FeatureVector {
-            values: vec![0.3; 56],
+            values: vec![0.3; ml_core::state_layout::STATE_DIM],
             timestamp: 1700000000,
         };
         let pred1 = adapter.predict(&fv).expect("first predict");