[indexing] Prevent integer overflow from large step values in C++ (#161707)

Fixes https://github.com/pytorch/pytorch/issues/160868
hmmm, I found an existing fix PR after I've finished this one. For reference, the old PR was
https://github.com/pytorch/pytorch/pull/147433/files.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/161707
Approved by: https://github.com/leslie-fang-intel, https://github.com/CaoE, https://github.com/mlazos
This commit is contained in:
thenumberouscode
2025-09-12 03:16:23 +00:00
committed by PyTorch MergeBot
parent 7eb92b076f
commit c140bf217f
3 changed files with 28 additions and 2 deletions

View File

@ -759,7 +759,8 @@ def slice_forward(
storage_offset = self.storage_offset() + start_val * strides[dim]
len = end_val - start_val
sizes[dim] = (len + step - 1) // step
# safely round-up for corresponding c++ impl
sizes[dim] = (len // step) + (1 if len % step != 0 else 0)
strides[dim] *= step
if self.is_quantized: