Assignment target is transposed when using jit.script and avanced indexing

🐛 Bug

When using advanced indexing for the target of an assignment myvar[scalar_tensor, slice, tensor, tensor] = ... the slice dimension is moved at the end if the function is jitted.

To Reproduce

import torch


def f(gt_batch_idx, gt_boxes):
    tgt_boxes = torch.zeros((1, 4, 64, 64))

    for i in range(gt_batch_idx.shape[0]):
        b = gt_batch_idx[i]  # works if using a simple int here

        mask_i, mask_j = torch.arange(10), torch.arange(10)

        print(tgt_boxes[b, :, mask_i, mask_j].shape)  # [10, 4] instead of [4, 10]
        tgt_boxes[b, :, mask_i, mask_j] = gt_boxes[i].view(4, 1)


f = torch.jit.script(f)  # works if commented

f(torch.tensor([0, 0], dtype=torch.long),
  torch.tensor([[31, 43, 58, 63], [22,  9, 45, 35]], dtype=torch.float))

Expected behavior

Same behaviour between torchscript and normal versions of the same code.

Environment

PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 10.2.0
Clang version: Could not collect
CMake version: version 3.19.3

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.7.1
[pip3] torchaudio==0.7.0a0+a853dff
[pip3] torchvision==0.8.2
[conda] blas                      1.0                         mkl  
[conda] cpuonly                   1.0                           0    pytorch
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.2.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.19.2           py38h54aff64_0  
[conda] numpy-base                1.19.2           py38hfa32c7d_0  
[conda] pytorch                   1.7.1               py3.8_cpu_0  [cpuonly]  pytorch
[conda] torchaudio                0.7.2                      py38    pytorch
[conda] torchvision               0.8.2                  py38_cpu  [cpuonly]  pytorch

cc @gmagogsfm