diff --git a/torch/csrc/jit/custom_operator.h b/torch/csrc/jit/custom_operator.h index 7dc3fcaee7f5..93a334ab9e08 100644 --- a/torch/csrc/jit/custom_operator.h +++ b/torch/csrc/jit/custom_operator.h @@ -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), options)); return *this; } + + template + RegisterOperators& op( + const std::string& name, + Implementation&& implementation) { + registrars_.emplace_back(std::make_shared(name, std::forward(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> registrars_; }; } // namespace jit