Caffe2 - C++ API
A deep learning, cross platform ML framework
replace_nan_op.h
1 #ifndef CAFFE_OPERATORS_REPLACE_NAN_OP_H_
2 #define CAFFE_OPERATORS_REPLACE_NAN_OP_H_
3 
4 #include "caffe2/core/context.h"
5 #include "caffe2/core/logging.h"
6 #include "caffe2/core/operator.h"
7 #include "caffe2/utils/math.h"
8 
9 namespace caffe2 {
10 
11 template <class Context>
12 class ReplaceNaNOp final : public Operator<Context> {
13  public:
14  USE_OPERATOR_CONTEXT_FUNCTIONS;
15  template <class... Args>
16  explicit ReplaceNaNOp(Args&&... args)
17  : Operator<Context>(std::forward<Args>(args)...) {}
18 
19  bool RunOnDevice() override {
20  return DispatchHelper<TensorTypes<float, double>>::call(this, Input(0));
21  }
22 
23  template <typename T>
24  void ReplaceNaN(const T& value, const int64_t size, const T* X, T* Y);
25 
26  template <typename T>
27  bool DoRunWithType() {
28  T value = this->template GetSingleArgument<T>("value", 0);
29 
30  auto& input = Input(0);
31 
32  auto* output = Output(0, input.sizes(), at::dtype<T>());
33 
34  const T* input_data = input.template data<T>();
35  T* output_data = output->template mutable_data<T>();
36 
37  ReplaceNaN<T>(value, input.numel(), input_data, output_data);
38 
39  return true;
40  }
41 };
42 
43 } // namespace caffe2
44 
45 #endif // CAFFE_OPERATORS_REPLACE_NAN_OP_H_
const Tensor & Input(int idx, DeviceType type=Context::GetDeviceType())
Retrieve a non-owning reference to the input at position &#39;idx&#39; for this operator. ...
Definition: operator.h:702
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13