1 #ifndef CAFFE2_OPERATORS_FLATTEN_OP_H_ 2 #define CAFFE2_OPERATORS_FLATTEN_OP_H_ 4 #include "caffe2/core/operator.h" 8 template <
class Context>
11 USE_OPERATOR_CONTEXT_FUNCTIONS;
15 axis_(OperatorBase::GetSingleArgument<int>(
"axis", 1)) {}
17 bool RunOnDevice()
override {
18 auto& input = Input(0);
19 auto* output = Output(0);
21 input.dims().size(), axis_,
"The rank of the tensor must be >= axis.");
22 output->Resize(input.size_to_dim(axis_), input.size_from_dim(axis_));
23 context_.template CopyItems<Context, Context>(
27 output->raw_mutable_data(input.meta()));
37 #endif // CAFFE2_OPERATORS_FLATTEN_OP_H_ Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...