Co-authored-by: Nikita Shulga <nshulga@meta.com>
Fixes https://github.com/pytorch/pytorch/issues/119054
resolved: https://github.com/pytorch/pytorch/pull/119750