Guard the boundary of index computed in compute_source_index_and_lambda (#89252)
Improve the fix in https://github.com/pytorch/pytorch/pull/89210
See discussion in https://github.com/pytorch/pytorch/issues/89212#issuecomment-1318911969
Pull Request resolved: https://github.com/pytorch/pytorch/pull/89252
Approved by: https://github.com/mingfeima, https://github.com/weiwangmeta
