FusedAdagrad.h 749 B

12345678910111213141516171819202122232425
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #include <ATen/core/Tensor.h>
  3. #include <ATen/native/DispatchStub.h>
  4. namespace at::native {
  5. using fused_adagrad_fn = void (*)(
  6. const at::Tensor& param,
  7. const at::Tensor& grad,
  8. const at::Tensor& state_sum,
  9. const at::Tensor& state_step,
  10. const double lr,
  11. const double lr_decay,
  12. const double weight_decay,
  13. const double eps,
  14. const bool maximize,
  15. const float* grad_scale_ptr);
  16. DECLARE_DISPATCH(fused_adagrad_fn, fused_adagrad_stub)
  17. } // namespace at::native
  18. #else
  19. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  20. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)