thrust_rl/utils/
cuda.rs

1//! Backend device helpers for the Burn training stack.
2//!
3//! After phase 5 of the Burn migration (#82), libtorch's CUDA runtime
4//! guard is gone — Burn picks its device from the backend type
5//! parameter (`B::Device`) instead of through a runtime probe. This
6//! module just exposes the default-device helper trainers use.
7
8/// Backend-generic default device helper for the Burn path.
9///
10/// Burn's device type is associated with the backend (`B::Device`); each
11/// backend implements `Default` for it, so for CPU (`NdArray`) this returns
12/// the singleton CPU device, for `Wgpu` the default GPU adapter, and so on.
13/// Trainer code that wants to stay backend-agnostic should call this
14/// instead of constructing a device manually.
15///
16/// # Example
17/// ```rust,no_run
18/// # #[cfg(feature = "training")]
19/// # {
20/// use burn::backend::NdArray;
21/// use thrust_rl::utils::cuda::default_burn_device;
22///
23/// let device = default_burn_device::<NdArray<f32>>();
24/// # let _ = device;
25/// # }
26/// ```
27#[cfg(feature = "training")]
28pub fn default_burn_device<B: burn::tensor::backend::Backend>()
29-> <B as burn::tensor::backend::BackendTypes>::Device {
30    <<B as burn::tensor::backend::BackendTypes>::Device as Default>::default()
31}
32
33#[cfg(test)]
34mod tests {
35    #[cfg(feature = "training")]
36    #[test]
37    fn test_default_burn_device_ndarray_smoke() {
38        // Sanity-check: the helper compiles, returns a device, and the
39        // returned device can actually allocate a tensor (i.e. it isn't
40        // some uninitialized placeholder).
41        use burn::{
42            backend::NdArray,
43            tensor::{Tensor, TensorData},
44        };
45
46        type B = NdArray<f32>;
47        let device = super::default_burn_device::<B>();
48        let t: Tensor<B, 1> =
49            Tensor::from_data(TensorData::new(vec![1.0_f32, 2.0, 3.0], [3]), &device);
50        assert_eq!(t.dims(), [3]);
51        let data: Vec<f32> = t.into_data().to_vec().unwrap();
52        assert_eq!(data, vec![1.0, 2.0, 3.0]);
53    }
54}
thrust_rl/utils/cuda.rs

thrust_rl/utils/
cuda.rs
