Program Listing for File step_lr.h

↰ Return to documentation for file (torch/csrc/api/include/torch/optim/schedulers/step_lr.h)

#pragma once

#include <torch/optim/schedulers/lr_scheduler.h>

namespace torch::optim {

class TORCH_API StepLR : public LRScheduler {
 public:
  StepLR(
      torch::optim::Optimizer& optimizer,
      const unsigned step_size,
      const double gamma = 0.1);

 private:
  std::vector<double> get_lrs() override;

  const unsigned step_size_;
  const double gamma_;
};
} // namespace torch::optim