pytorch

bbf91554 - Cherry-pick of #97214 (#98718)

Cherry-pick of #97214 (#98718) * merge * add missing import * better merge conflict resolution * cherry pick not clean --------- Co-authored-by: Michael Gschwind <mikekg@meta.com>