enhance next_power_of_2 function (#120153)
Fixes #120152
cc @ezyang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @jansel
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120153
Approved by: https://github.com/jansel
