Spaces:
Running
Running
File size: 647 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
#pragma once
#include <ATen/native/DispatchStub.h>
#include <ATen/core/ATen_fwd.h>
namespace at {
class Tensor;
namespace native {
using _amp_foreach_non_finite_check_and_unscale_cpu__fn = void (*)(
TensorList,
Tensor&,
const Tensor&);
using _amp_update_scale_cpu__fn = Tensor& (*)(
Tensor&,
Tensor&,
const Tensor&,
double,
double,
int64_t);
DECLARE_DISPATCH(_amp_foreach_non_finite_check_and_unscale_cpu__fn, _amp_foreach_non_finite_check_and_unscale_cpu_stub);
DECLARE_DISPATCH(_amp_update_scale_cpu__fn, _amp_update_scale_cpu_stub);
} // namespace native
} // namespace at
|