burn_train/learner/
regression.rs

1use crate::metric::processor::ItemLazy;
2use crate::metric::{Adaptor, LossInput};
3use burn_core::tensor::Tensor;
4use burn_std::ExecutionError;
5
6/// Regression output adapted for the loss metric.
7#[derive(new)]
8pub struct RegressionOutput {
9    /// The loss.
10    pub loss: Tensor<1>,
11
12    /// The predicted values. Shape: \[batch_size, num_targets\].
13    pub output: Tensor<2>,
14
15    /// The ground truth values. Shape: \[batch_size, num_targets\].
16    pub targets: Tensor<2>,
17}
18
19impl Adaptor<LossInput> for RegressionOutput {
20    fn adapt(&self) -> LossInput {
21        LossInput::new(self.loss.clone())
22    }
23}
24
25impl ItemLazy for RegressionOutput {
26    fn sync(self) -> Result<Self, ExecutionError> {
27        // No readback: the metrics compute on the device the tensors live on
28        // and read back only their final scalars. Flushing dispatches the
29        // producing stream's buffered work so the metric thread doesn't wait
30        // on an idle queue; a training item's float tensors come off the
31        // autodiff backend entirely, so the metric thread neither retains the
32        // tape nor pays its dispatch.
33        self.loss.device().flush()?;
34
35        Ok(RegressionOutput {
36            output: self.output.without_autodiff(),
37            loss: self.loss.without_autodiff(),
38            targets: self.targets.without_autodiff(),
39        })
40    }
41}
burn_train/learner/regression.rs

burn_train/learner/
regression.rs
