torch::jit::RegisterOperators forwards to c10::RegisterOperators

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/20383

Reviewed By: zdevito

Differential Revision: D15300937

fbshipit-source-id: 740fe323fc0945759651116ae61aff4d36319d73
This commit is contained in:
Sebastian Messmer
2019-05-20 12:37:54 -07:00
committed by Facebook Github Bot
parent 74bdcd44c4
commit 9fbce974c9

View File

@ -246,11 +246,29 @@ struct TORCH_API RegisterOperators {
RegisterOperators& op(
const std::string& name,
Implementation&& implementation,
OperatorOptions options = OperatorOptions()) {
OperatorOptions options) {
registerOperator(createOperator(
name, std::forward<Implementation>(implementation), options));
return *this;
}
template <typename Implementation>
RegisterOperators& op(
const std::string& name,
Implementation&& implementation) {
registrars_.emplace_back(std::make_shared<c10::RegisterOperators>(name, std::forward<Implementation>(implementation)));
return *this;
}
private:
// A c10::RegisterOperators instance is not copyable, so to make
// torch::jit::RegisterOperators copyable, we use shared_ptrs.
// We need to keep the c10::RegisterOperators instances around
// because this is an RAII pattern. In the destructor, the registered
// ops get de-registered.
std::vector<std::shared_ptr<c10::RegisterOperators>> registrars_;
};
} // namespace jit