Module state 

Traits

OptimizerState

OptimizerState trait defines the interface for managing the state of an optimizer. It provides methods to initialize the optimizer state and update the model weights. This trait is useful to implement custom parameter update rules, as used in different optimization algorithms, like RMSProp, Adam, etc.