Refactor lambda post hook. (#37025)

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

This allows us to reuse this framework in other places.

Test Plan:
buck test mode/dev-nosan
caffe2/torch/fb/distributed/model_parallel/tests:test_dist_optim --
test_optimizer_hook

Differential Revision: D20958327

fbshipit-source-id: 2a37dae3687fea8820427e174900111b58673194

torch/csrc/autograd/utils/lambda_post_hook.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <torch/csrc/autograd/function_hook.h>
+
+namespace torch {
+namespace autograd {
+namespace utils {
+
+// Turns lambda into a torch::autograd::FunctionPostHook.
+class LambdaPostHook : public torch::autograd::FunctionPostHook {
+  using variable_list = std::vector<torch::autograd::Variable>;
+
+ public:
+  // The lambda function takes as arguments the outputs and inputs of the
+  // autograd function and can modify the outputs of the autograd function by
+  // returning a new output if needed.
+  /* implicit */ LambdaPostHook(
+      std::function<variable_list(const variable_list&, const variable_list&)>
+          fn)
+      : fn_(std::move(fn)) {}
+
+  variable_list operator()(
+      const variable_list& outputs,
+      const variable_list& inputs) override {
+    return fn_(outputs, inputs);
+  }
+
+ protected:
+  std::function<variable_list(const variable_list&, const variable_list&)> fn_;
+};
+
+} // namespace utils
+} // namespace autograd
+} // namespace torch