1 #ifndef CAFFE2_OPERATORS_TOP_K_H_ 2 #define CAFFE2_OPERATORS_TOP_K_H_ 4 #include "caffe2/core/logging.h" 5 #include "caffe2/core/operator.h" 6 #include "caffe2/utils/math.h" 10 template <
typename T,
class Context>
13 USE_OPERATOR_CONTEXT_FUNCTIONS;
17 OP_SINGLE_ARG(
int,
"k", k_, -1),
18 OP_SINGLE_ARG(
int,
"axis", axis_, -1) {
19 CAFFE_ENFORCE(k_ >= 1,
"k argument must be >= 1");
24 bool RunOnDevice()
override;
31 template <
typename T,
class Context>
34 USE_OPERATOR_CONTEXT_FUNCTIONS;
38 OP_SINGLE_ARG(
int,
"axis", axis_, -1) {}
42 bool RunOnDevice()
override;
50 #endif // CAFFE2_OPERATORS_TOP_K_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 ...