17 #include "caffe2/operators/stump_func_op.h" 22 bool StumpFuncOp<float, float, CPUContext>::RunOnDevice() {
24 const float* in_data = in.template data<float>();
25 auto* out = Output(0);
27 float* out_data = out->template mutable_data<float>();
28 for (
int i = 0; i < in.size(); i++) {
29 out_data[i] = (in_data[i] <= threshold_) ? low_value_ : high_value_;
34 REGISTER_CPU_OPERATOR(StumpFunc, StumpFuncOp<float, float, CPUContext>);
36 OPERATOR_SCHEMA(StumpFunc)
39 .Input(0,
"X",
"tensor of float")
40 .Output(0,
"Y",
"tensor of float")
42 Converts each input element into either high_ or low_value 43 based on the given threshold. 46 NO_GRADIENT(StumpFunc); A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...