1 #pragma once 2 3 #include <torch/optim/schedulers/lr_scheduler.h> 4 5 namespace torch { 6 namespace optim { 7 8 class TORCH_API StepLR : public LRScheduler { 9 public: 10 StepLR( 11 torch::optim::Optimizer& optimizer, 12 const unsigned step_size, 13 const double gamma = 0.1); 14 15 private: 16 std::vector<double> get_lrs() override; 17 18 const unsigned step_size_; 19 const double gamma_; 20 }; 21 } // namespace optim 22 } // namespace torch 23