Caffe2 - C++ API
A deep learning, cross platform ML framework
enforce_finite_cpu.cc
1 #include <ATen/core/dispatch/KernelRegistration.h>
2 #include "caffe2/operators/experimental/c10/schemas/enforce_finite.h"
3 #include "caffe2/utils/math.h"
4 #include "caffe2/core/tensor.h"
5 
7 using caffe2::Tensor;
8 
9 namespace caffe2 {
10 namespace {
11 template <class DataType>
12 void enforce_finite_op_impl_cpu(const at::Tensor& input_) {
13  Tensor input{C10Tensor(input_)};
14  const DataType* input_data = input.template data<DataType>();
15  auto size = input.numel();
16 
17  for (auto i = 0; i < size; i++) {
18  CAFFE_ENFORCE(
19  std::isfinite(input_data[i]),
20  "Index ",
21  i,
22  " is not finite (e.g., NaN, Inf): ",
23  input_data[i]);
24  }
25 }
26 } // namespace
27 } // namespace caffe2
28 
29 namespace c10 {
30 C10_REGISTER_KERNEL(caffe2::ops::EnforceFinite)
31  .kernel<decltype(caffe2::enforce_finite_op_impl_cpu<float>), &caffe2::enforce_finite_op_impl_cpu<float>>()
32  .dispatchKey(CPUTensorId());
33 } // namespace c10
Tensor class holds a shared pointer to the implementation TensorImpl, redirects API calls to TensorIm...
Definition: tensor.h:25
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Definition: context.h:40
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
To register your own kernel for an operator, do in one (!) cpp file: C10_REGISTER_KERNEL(OperatorHand...
Definition: alias_info.h:7