huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 11:08:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Set the learning rate. (#687)

Browse Source
This commit is contained in:
Laurent Mazare
2023-08-31 09:03:40 +02:00
committed by GitHub
parent 8e84d8a59b
commit d210c71d77
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-nn/src/optim.rs
View File

@ -58,6 +58,10 @@ impl SGD {
let grads = loss.backward()?; let grads = loss.backward()?;
self.step(&grads) self.step(&grads)
} }
pub fn set_learning_rate(&mut self, lr: f64) {
self.learning_rate = lr
}
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@ -127,6 +131,10 @@ impl AdamW {
Self::new(vars, params) Self::new(vars, params)
} }
pub fn set_learning_rate(&mut self, lr: f64) {
self.params.lr = lr
}
pub fn step(&mut self, grads: &candle::backprop::GradStore) -> Result<()> { pub fn step(&mut self, grads: &candle::backprop::GradStore) -> Result<()> {
self.step_t += 1; self.step_t += 1;
let lr = self.params.lr; let lr = self.params.lr;