Fixes #120788

Pull Request resolved: https://github.com/pytorch/pytorch/pull/121618
Approved by: https://github.com/xuhancn, https://github.com/jgong5, https://github.com/mingfeima, https://github.com/seemethere

(cherry picked from commit 03717430cc54609189cc7df593b2c96a99fb7f55)

Co-authored-by: CaoE <e.cao@intel.com>