| // Copyright 2004-present Facebook. All Rights Reserved. | |
| #ifndef CONDITIONAL_OP_H | |
| #define CONDITIONAL_OP_H | |
| #include "caffe2/core/context.h" | |
| #include "caffe2/core/operator.h" | |
| #include "caffe2/core/tensor.h" | |
| namespace caffe2 { | |
| template <class Context> | |
| class ConditionalOp final : public Operator<Context> { | |
| public: | |
| USE_OPERATOR_CONTEXT_FUNCTIONS; | |
| ConditionalOp(const OperatorDef& operator_def, Workspace* ws) | |
| : Operator<Context>(operator_def, ws) {} | |
| bool RunOnDevice() override; | |
| }; | |
| } // caffe2 | |
| #endif |