pub fn with_grad<T, F>(f: F) -> T where F: FnOnce() -> T,
Runs a closure explicitly keeping track of gradients, this could be run within a no_grad closure for example.