Caffe2 - C++ API
A deep learning, cross platform ML framework
replace_nan_op.cc
1 
17 #include "caffe2/operators/replace_nan_op.h"
18 
19 namespace caffe2 {
20 
21 template <>
22 template <typename T>
23 void ReplaceNaNOp<CPUContext>::ReplaceNaN(
24  const T& value,
25  const TIndex size,
26  const T* X,
27  T* Y) {
28  for (TIndex i = 0; i < size; i++) {
29  if (std::isnan(X[i])) {
30  Y[i] = value;
31  } else {
32  Y[i] = X[i];
33  }
34  }
35 }
36 
37 REGISTER_CPU_OPERATOR(ReplaceNaN, ReplaceNaNOp<CPUContext>);
38 
39 OPERATOR_SCHEMA(ReplaceNaN)
40  .NumInputs(1)
41  .NumOutputs(1)
42  .AllowInplace({{0, 0}})
43  .IdenticalTypeAndShape()
44  .SetDoc(R"DOC(
45 Replace the NaN (not a number) element in the input tensor with argument `value`
46 )DOC")
47  .Arg("value (optional)", "the value to replace NaN, the default is 0")
48  .Input(0, "input", "Input tensor")
49  .Input(1, "output", "Output tensor");
50 
51 SHOULD_NOT_DO_GRADIENT(ReplaceNaN);
52 
53 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.