hodu_nn/losses/
mse.rs

1use crate::compat::*;
2use crate::module::Module;
3use hodu_core::{error::HoduResult, scalar::Scalar, tensor::Tensor};
4
5#[derive(Module, Clone, Default)]
6#[module(inputs = 2)]
7pub struct MSELoss;
8
9impl MSELoss {
10    pub fn new() -> Self {
11        Self
12    }
13
14    pub fn forward(&self, (pred, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
15        let diff = pred.sub(target)?;
16        let squared = diff.pow_scalar(Scalar::from_f32(2.0, diff.dtype()))?;
17
18        let squared_size = squared.layout().size();
19        let num_elements = Scalar::new(squared_size).to_dtype(squared.dtype());
20        let mean = squared.sum_all()?.div_scalar(num_elements)?;
21
22        Ok(mean)
23    }
24}
hodu_nn/losses/mse.rs

hodu_nn/losses/
mse.rs
