Update amp custom_fwd, custom_bwd usage for torch 2.4.0 compatibility#54

Merged

yzhangcs merged 1 commit intofla-org:mainfrom mirceamironenco:fix-torch2.4.0-amp-futurewarning

Aug 25, 2024

+60-67

Commits on Aug 25, 2024

Update amp custom_fwd, custom_bwd usage for torch 2.4.0 compatibility
mirceamironenco
committed