DeepSpeed

6ea345ae - Simplify leaf module hook (#7592)

Simplify leaf module hook (#7592) This PR simplifies hooks for leaf module using PyTorch's API. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>