Struct burn_core::optim::GradientsParams

pub struct GradientsParams { /* private fields */ }

Data type that contains gradients for parameters.

Implementations

impl GradientsParams

pub fn new() -> Self

pub fn get<B, const D: usize>(&self, id: &ParamId) -> Option<Tensor<B, D>>where B: Backend,

Get the gradients for the given parameter id.

Notes

You should use remove if you want to get the gradients only one time.

pub fn remove<B, const D: usize>(&self, id: &ParamId) -> Option<Tensor<B, D>>where B: Backend,

Remove the gradients for the given parameter id.

pub fn register<B, const D: usize>(&mut self, id: ParamId, value: Tensor<B, D>)where B: Backend,

Register a gradients tensor for the given parameter id.

Notes

If a tensor is already registered for the given parameter id, it will be replaced.

pub fn len(&self) -> usize

The number of gradients tensors registered.

pub fn is_empty(&self) -> bool

If any tensor is contained.

pub fn to_device<B: ADBackend, M: ADModule<B>>( self, device: &B::Device, module: &M ) -> Self

Change the device of each tensor gradients registered for the given module.

pub fn from_grads<B: ADBackend, M: ADModule<B>>( grads: B::Gradients, module: &M ) -> Self

Extract each tensor gradients for the given module.

Trait Implementations

impl Default for GradientsParams

fn default() -> GradientsParams

Returns the “default value” for a type.