mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
71ca600af9
Summary: Since caffe2 and torch have been consolidated, CAFFE2_API should be merged with TORCH_API. Addresses a TODO. Manually edited some references of the removed `CAFFE2_API`: * `CONTRIBUTING.md` * `caffe2/proto/CMakeLists.txt` * `cmake/ProtoBuf.cmake` * `c10/macros/Export.h` * `torch/csrc/WindowsTorchApiMacro.h` Pull Request resolved: https://github.com/pytorch/pytorch/pull/49496 Reviewed By: malfet, samestep Differential Revision: D25600726 Pulled By: janeyx99 fbshipit-source-id: 7e068d959e397ac183c097d7e9a9afeca5ddd782
25 lines
678 B
C++
25 lines
678 B
C++
#pragma once
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/proto/caffe2_pb.h"
|
|
#include "caffe2/transforms/single_op_transform.h"
|
|
#include "caffe2/utils/proto_utils.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
class TORCH_API ConvToNNPackTransform : public SingleOpTransform {
|
|
protected:
|
|
// Specify what the op needs to be to match the pattern.
|
|
bool MatchOperator(const OperatorDef& op) override {
|
|
return (
|
|
op.type() == "Conv" && op.device_option().device_type() == PROTO_CPU &&
|
|
op.engine() != "NNPACK");
|
|
}
|
|
|
|
// Specify how the operator should be replaced.
|
|
void ReplaceOperator(OperatorDef* op) override {
|
|
op->set_engine("NNPACK");
|
|
}
|
|
};
|
|
|
|
} // namespace caffe2
|