foxhunt/ml/tests/verify_dqn_cuda.rs

#[cfg(test)]
mod verify_dqn_cuda_tests {
    use candle_core::{DType, Device, Tensor};
    use ml::dqn::{WorkingDQN, WorkingDQNConfig};

    #[test]
    fn test_dqn_uses_cuda_device() -> anyhow::Result<()> {
        // Create DQN with default config
        let config = WorkingDQNConfig::emergency_safe_defaults();
        let dqn = WorkingDQN::new(config.clone())?;

        // Create test input on CPU first
        let state_cpu = Tensor::zeros(&[1, config.state_dim], DType::F32, &Device::Cpu)?;

        // Forward pass
        let output = dqn.forward(&state_cpu)?;

        // Check output device
        println!("Output tensor device: {:?}", output.device());
        println!("Is CUDA: {}", output.device().is_cuda());
        println!("Is CPU: {}", output.device().is_cpu());

        // The output should be on CUDA if GPU is available
        if cfg!(feature = "cuda") {
            assert!(
                output.device().is_cuda(),
                "DQN should use CUDA device when available"
            );
            println!("✅ DQN is using CUDA GPU acceleration");
        } else {
            println!("⚠️  CUDA feature not enabled, using CPU");
        }

        Ok(())
    }

    #[test]
    fn test_device_selection() -> anyhow::Result<()> {
        let device = Device::cuda_if_available(0)?;
        println!("Selected device: {:?}", device);
        println!("Is CUDA: {}", device.is_cuda());

        if device.is_cuda() {
            println!("✅ CUDA device available");

            // Try allocating a small tensor on GPU
            let test_tensor = Tensor::zeros(&[100, 100], DType::F32, &device)?;
            println!("Test tensor shape: {:?}", test_tensor.shape());
            println!("Test tensor device: {:?}", test_tensor.device());
        } else {
            println!("⚠️  Falling back to CPU");
        }

        Ok(())
    }
}