Repository: pytorch
Issue: Issue with `upsample_nearest2d` decomposition
Labels: triaged,module: decompositions

### 🐛 Describe the bug

After the decomposition of `upsample_nearest2d` getting updated here: https://github.com/pytorch/pytorch/commit/8ecda19607a819d0387341b6c2d31d29e05a33ab, it results out in invalid index accesses. For example. consider the following case, where both the PyTorch implementation of the op and decomp are used:
```
import torch
from typing import Callable, cast, Iterable, List, Optional, Tuple, Union
from torch import sym_float, sym_int, Tensor
import torch._prims_common as utils

def _compute_upsample_nearest_indices(input, output_size, scales):
    # For each dim in output_size, compute the set of input indices used
    # to produce the upsampled output.
    indices = []
    num_spatial_dims = len(output_size)
    for d in range(num_spatial_dims):
        # Math matches aten/src/ATen/native/cpu/UpSampleKernel.cpp
        # Indices are computed as following:
        # scale = isize / osize
        # input_index = floor(output_index * scale)
        # Same as OpenCV INTER_NEAREST
        osize = output_size[d]
        output_indices = torch.arange(osize, dtype=input.dtype, device=input.device)
        isize = input.shape[-num_spatial_dims + d]
        scale = isize / (isize * scales[d]) if scales[d] is not None else isize / osize
        input_indices = (output_indices * scale).to(torch.int64)
        for _ in range(num_spatial_dims - 1 - d):
            input_indices = input_indices.unsqueeze(-1)
        indices.append(input_indices)
    return tuple(indices)

def upsample_nearest2d(
    input: Tensor,
    output_size: List[int],
    scales_h: Optional[float] = None,
    scales_w: Optional[float] = None,
) -> Tensor:
    h_indices, w_indices = _compute_upsample_nearest_indices(
        input, output_size, (scales_h, scales_w)
    )
    result = input[:, :, h_indices, w_indices]

    # convert output to correct memory format, if necessary
    memory_format = utils.suggest_memory_format(input)

    # following "heuristic: only use channels_last path when it's faster than the contiguous path"
    _, n_channels, _, _ = input.shape
    if input.device.type == "cuda" and n_channels < 4:
        memory_format = torch.contiguous_format

    result = result.contiguous(memory_format=memory_format)

    return result

inputVec = torch.rand((2, 3, 2, 2))
out_torch = torch._C._nn.upsample_nearest2d(
    inputVec, output_size=[6, 10], scales_h=2.3, scales_w=4.7
)

out_decomp = upsample_nearest2d(
    inputVec, output_size=[6, 10], scales_h=2.3, scales_w=4.7
)

print(out_torch == out_decomp)
```
If you run this code the decomp will give the following error:
```
Traceback (most recent call last):
  File "/home/vivek/work/misc/vivekkhandelwal1-pytorch/upsample_nearest2d.py", line 58, in <module>
    out_decomp = upsample_nearest2d(
  File "/home/vivek/work/misc/vivekkhandelwal1-pytorch/upsample_nearest2d.py", line 38, in upsample_nearest2d
    result = input[:, :, h_indices, w_indices]
IndexError: index 2 is out of bounds for dimension 0 with size 2
```
If we revert the changes made at this line (https://github.com/pytorch/pytorch/blob/8ecda19607a819d0387341b6c2d31d29e05a33ab/torch/_decomp/decompositions.py#L2059) particularly, then it works fine.

### Versions

PyTorch version: 2.0.0a0+gitf58ba55
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.10.9 (main, Dec  7 2022, 01:12:00) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.6.55
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 470.161.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.1
[pip3] torch==2.0.0a0+gitf58ba55
[conda] Could not collect

cc @SherlockNoMad @ngimel