#ifndef CAFFE_OPERATORS_REPLACE_NAN_OP_H_ #define CAFFE_OPERATORS_REPLACE_NAN_OP_H_ #include "caffe2/core/context.h" #include "caffe2/core/logging.h" #include "caffe2/core/operator.h" #include "caffe2/utils/math.h" namespace caffe2 { template class ReplaceNaNOp final : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; template explicit ReplaceNaNOp(Args&&... args) : Operator(std::forward(args)...) {} bool RunOnDevice() override { return DispatchHelper>::call(this, Input(0)); } template void ReplaceNaN(const T& value, const int64_t size, const T* X, T* Y); template bool DoRunWithType() { T value = this->template GetSingleArgument("value", 0); auto& input = Input(0); auto* output = Output(0, input.sizes(), at::dtype()); const T* input_data = input.template data(); T* output_data = output->template mutable_data(); ReplaceNaN(value, input.numel(), input_data, output_data); return true; } }; } // namespace caffe2 #endif // CAFFE_OPERATORS_REPLACE_NAN_OP_H_