Function with_device 

pub fn with_device<F, R>(device: Device, f: F) -> R
where
    F: FnOnce() -> R,

Execute a closure with a specific device context

This function temporarily switches to the specified device for the duration of the closure, then automatically restores the previous device. This is the recommended way to execute code with a specific device context.

Arguments

• device - The device to use for the closure
• f - The closure to execute

Returns

The result of the closure

Thread Safety

This function is thread-safe and only affects the current thread’s context.

Examples

use train_station::{Device, with_device, current_device};

let original_device = current_device();

let result = with_device(Device::cpu(), || {
    assert_eq!(current_device(), Device::cpu());
    // Perform operations with CPU device
    42
});

assert_eq!(result, 42);
assert_eq!(current_device(), original_device);