Function batch_norm_train 

pub fn batch_norm_train<const D: usize>(
    input: Tensor<D>,
    gamma: Tensor<1>,
    beta: Tensor<1>,
    epsilon: f64,
) -> BatchNormTrainOutput<D>

Applies batch normalization using statistics computed from the input batch.

input has shape [batch, channels, ...]; gamma and beta have shape [channels].

Returns the normalized input with its original shape, along with the batch mean and biased (population) variance, both with shape [channels]. Statistics are computed over every dimension except the channel dimension. The returned variance excludes epsilon.

This function does not update running statistics. For normalization using explicitly supplied statistics, use batch_norm.

Autodiff

This function can be used with or without autodiff. It does not enable gradient tracking.

On an autodiff-enabled device, gradients through the normalized output account for the dependence of the batch statistics on the input. The returned mean and variance remain on the same device but are detached.