Fix https://github.com/pytorch/pytorch/issues/108978

Cherry-pick of https://github.com/pytorch/pytorch/pull/109557 into release/2.1 branch 


(cherry picked from commit 00871189972e81a5fde230bc08137be14c59f178)

Co-authored-by: Li-Huai (Allan) Lin <qqaatw@gmail.com>