FusedSGD.h 770 B

1234567891011121314151617181920212223242526
  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_sgd_fn = void (*)(
  6. const at::Tensor& param,
  7. const at::Tensor& grad,
  8. const at::Tensor& momentum_buffer,
  9. const double weight_decay,
  10. const double momentum,
  11. const double lr,
  12. const double dampening,
  13. const bool nesterov,
  14. const bool maximize,
  15. const bool is_first_step,
  16. const float* grad_scale_ptr);
  17. DECLARE_DISPATCH(fused_sgd_fn, fused_sgd_stub)
  18. } // namespace at::native
  19. #else
  20. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  21. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)