Feature request - allow reuse of Gradients
Author: fuzingCreated Sep 15, 2026Updated Sep 17, 2026
It would be useful to be able to reuse Gradients - For example (plain step function):
fn step(&self, item: FlowRegFormerTrainingBatch) -> TrainOutput<RegressionOutput> {
// Run forward pass, calculate gradients and return them along with the output
let item = self.forward(item);
let grads = item.loss.backward();
TrainOutput::new(self, grads, item)
}I'm attempting to instrument/log gradient magnitudes, and find something like this useful:
fn step(&self, item: FlowRegFormerTrainingBatch) -> TrainOutput<RegressionOutput> {
// Run forward pass, calculate gradients and return them along with the output
let item = self.forward(item);
let grads = item.loss.backward();
// gradient examination
let model = self;
let grads_params = GradientsParams::from_grads(grads, model);
let mut inspector = GradientInspector::new(&grads_params, true);
model.visit(&mut inspector);
panic!("completed");
// end gradient examination
TrainOutput::new(self, grads, item)
}The problem is that GradientsParams::from_grads() consumes the gradients, which subsequently need to be returned from the step method in the TrainOutput struct.
My current workaround is to panic! after examining the gradients, meaning I can only examine them once at the start of an epoch (i.e. stop training then restart from last epoch -> examine gradients -> panic!).
It would be useful if the Gradients were clonable/copyable/reusable, or if GradientParams::from_grads() took a reference to the generated gradients and didn't consume them.
Thank you
Source: tracel-ai/burn