Repository: pytorch
Issue: torch.jit.script mean(keepdim=True) segfaults on GPU
Labels: oncall: jit

### 🐛 Describe the bug

When wrapping a `torch.nn.Module` with `torch.jit.script`, calling `torch.mean` with `keepdim=True` leads to a crash in backpropagation. Calling `torch.mean` followed by `unsqueeze` does not.

Crash message:
`RuntimeError: invalid vector subscript`

Reproducing code
```
#!/usr/bin/env python
import torch

N_Floats=50
Batch_Size = 8
Linear_Width = 512
N_Actions = 12

class Agent(torch.nn.Module):
	def __init__(self,Linear_Width:int):
		super().__init__()
		self.A_head = torch.nn.Linear(N_Floats,N_Actions)
	def forward(self, inputs):
		A = self.A_head(inputs)
		#Q = V + A - A.mean(dim=-1).unsqueeze(-1) #Works
		Q = A - A.mean(dim=-1,keepdim=True) #Crashes with "invalid vector subscript"
		return Q

def learn_on_batch(model,optimizer):
	Action_Indexes = torch.randint(low=0,high=N_Actions,size=(Batch_Size,1),device='cuda',dtype=torch.int64)
	Inputs = torch.rand(size=(Batch_Size,N_Floats),device='cuda',dtype=torch.float32)
	target = torch.ones(Batch_Size,device='cuda')
	outputs = model(Inputs)
	outputs = torch.gather(outputs, 1, Action_Indexes.type(torch.int64))
	loss = torch.nn.functional.mse_loss(outputs,target)
	total_loss = torch.sum(loss)
	optimizer.zero_grad(set_to_none=True)
	total_loss.backward()
	optimizer.step()

if __name__ == '__main__':
	model =  torch.jit.script(Agent(Linear_Width)).to("cuda")
	optimizer = torch.optim.Adam(model.parameters())
	while True:
		learn_on_batch(model,optimizer)
```
		
		



### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.10.11 | packaged by conda-forge | (main, May 10 2023, 18:51:25) [MSC v.1934 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 12.0.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4070
Nvidia driver version: 535.98
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3401
DeviceID=CPU0
Family=107
L2CacheSize=8192
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=3401
Name=AMD Ryzen 9 5950X 16-Core Processor
ProcessorType=3
Revision=8448

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.0.1
[pip3] torch-tb-profiler==0.4.1
[pip3] torchaudio==2.0.2
[pip3] torchrl==0.1.1
[conda] blas                      2.117                       mkl    conda-forge
[conda] blas-devel                3.9.0              17_win64_mkl    conda-forge
[conda] libblas                   3.9.0              17_win64_mkl    conda-forge
[conda] libcblas                  3.9.0              17_win64_mkl    conda-forge
[conda] liblapack                 3.9.0              17_win64_mkl    conda-forge
[conda] liblapacke                3.9.0              17_win64_mkl    conda-forge
[conda] mkl                       2022.1.0           h6a75c08_874    conda-forge
[conda] mkl-devel                 2022.1.0           h57928b3_875    conda-forge
[conda] mkl-include               2022.1.0           h6a75c08_874    conda-forge
[conda] numpy                     1.24.3          py310hd02465a_0    conda-forge
[conda] pytorch                   2.0.1           py3.10_cuda11.8_cudnn8_0    pytorch
[conda] pytorch-cuda              11.8                 h24eeafa_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.4.1                    pypi_0    pypi
[conda] torchaudio                2.0.2                    pypi_0    pypi
[conda] torchrl                   0.1.1                    pypi_0    pypi

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel