1 #ifndef CAFFE2_OPERATORS_INT8_FLATTEN_OP_H_ 2 #define CAFFE2_OPERATORS_INT8_FLATTEN_OP_H_ 4 #include "caffe2/core/context.h" 5 #include "caffe2/core/operator.h" 6 #include "caffe2/core/tensor_int8.h" 7 #include "caffe2/operators/quantized/int8_utils.h" 15 template <
class... Args>
18 axis_(this->
template GetSingleArgument<int>(
"axis", 1)) {}
20 bool RunOnDevice()
override {
23 int32_t Y_offset = this->
template GetSingleArgument<int>(
"Y_zero_point", 0);
24 auto Y_scale = this->
template GetSingleArgument<float>(
"Y_scale", 1);
25 CHECK_EQ(Y_offset, X.zero_point);
26 CHECK_EQ(Y_scale, X.scale);
28 Y->zero_point = Y_offset;
30 X.t.sizes().size(), axis_,
"The rank of the tensor must be >= axis.");
31 Y->t.Resize(X.t.size_to_dim(axis_), X.t.size_from_dim(axis_));
32 context_.CopyItemsToCPU(
36 Y->t.raw_mutable_data(X.t.dtype()));
48 #endif // CAFFE2_OPERATORS_INT8_FLATTEN_OP_H_
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...