mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
6b656565ab
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/24393 Ability to register hook on a variable, similar to python autograd API. register_hook will take a function as argument and create a CppFunctionPreHook similar to PyFunctionPreHook. It will return the index of the hook which can be passed to remove_hook to disable the hook. Test Plan: Added tests. Differential Revision: D16861722 fbshipit-source-id: d08047f932e38c7bde04283a18b2d0311c8ad604
18 lines
506 B
C++
18 lines
506 B
C++
#pragma once
|
|
#include <torch/csrc/autograd/function_hook.h>
|
|
#include <functional>
|
|
#include <memory>
|
|
|
|
namespace torch { namespace autograd {
|
|
|
|
using hooks_list = std::vector<std::function<Variable(const Variable&)>>;
|
|
|
|
struct CppFunctionPreHook : public FunctionPreHook {
|
|
CppFunctionPreHook(const std::shared_ptr<hooks_list> &hooks, int value_idx);
|
|
variable_list operator()(const variable_list& values) override;
|
|
|
|
std::shared_ptr<hooks_list> hooks_;
|
|
int value_idx_;
|
|
};
|
|
}} // namespace torch::autograd
|