Skip to main content

with_stream

Function with_stream 

Source
pub fn with_stream<R>(stream: impl Into<Stream>, f: impl FnOnce() -> R) -> R
Expand description

Runs f with stream as this thread’s default stream, and restores the previous default afterwards, also when f panics. A device stands for its first stream.

This replaces per-operation stream arguments: every operation f records goes on stream.

use nuvec_core::{Device, array, default_stream, with_stream};

let x = array![1.0, 2.0];
let y = with_stream(Device::cpu(), || x.exp().sum(..));
assert_eq!(y.device(), Device::cpu());
assert_eq!(default_stream(), nuvec_core::default_stream()); // restored