Add the optimizer trait. (#702)

This commit is contained in:
Laurent Mazare
2023-09-01 13:55:39 +02:00
committed by GitHub
parent f2d476ca65
commit 7529531056
6 changed files with 69 additions and 54 deletions

View File

@ -24,7 +24,7 @@ pub use init::Init;
pub use layer_norm::{layer_norm, rms_norm, LayerNorm, LayerNormConfig, RmsNorm};
pub use linear::{linear, linear_no_bias, Linear};
pub use ops::Dropout;
pub use optim::{AdamW, ParamsAdamW, SGD};
pub use optim::{AdamW, Optimizer, ParamsAdamW, SGD};
pub use rnn::{gru, lstm, GRUConfig, LSTMConfig, GRU, LSTM, RNN};
pub use var_builder::VarBuilder;
pub use var_map::VarMap;