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
