Set the learning rate. (#687)

candle-nn/src/optim.rs

@@ -58,6 +58,10 @@ impl SGD {
         let grads = loss.backward()?;
         self.step(&grads)
     }
+
+    pub fn set_learning_rate(&mut self, lr: f64) {
+        self.learning_rate = lr
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -127,6 +131,10 @@ impl AdamW {
         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<()> {
         self.step_t += 1;
         let lr = self.params.lr;