Repository: pytorch
Issue: The reload `MultiLabelMarginLoss` will have different gradients on cuda
Labels: oncall: jit

### 🐛 Describe the bug

The reload `MultiLabelMarginLoss` will have different gradients on cuda

```py
import torch

device = 'cuda'
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.arg_class = torch.nn.MultiLabelMarginLoss()
        self.input_1 = torch.tensor([[15, 0, 0, 0]], dtype=torch.int64, device=device)

    def forward(self, input_0):
        arg_class = self.arg_class
        input_1 = self.input_1
        return arg_class(input_0, input_1)

input = torch.tensor([[-0.2898, -0.9470, -0.4042, -0.6775]], dtype=torch.float32, device=device)

m = M().to(device)
jit_m = torch.jit.trace(m, (input.clone(), ))

input1 = input.clone().requires_grad_()
jit_m(input.clone())
jit_m(input1).sum().backward()
print(input1.grad)

import io
buffer = io.BytesIO()
torch.jit.save(jit_m, buffer)
buffer.seek(0)
new_m = torch.jit.load(buffer)

input2 = input.clone().requires_grad_()
new_m(input2).sum().backward()
print(input2.grad)
```

```
tensor([[-2.2500,  1.0000,  1.0000,  1.0000]], device='cuda:0')
tensor([[-2.2500,  0.7500,  0.7500,  0.7500]], device='cuda:0')
```

This is caused by the out-of-bound index `15`. By contrast, on cpu, it will just raise an error that the target is out of bound

### Versions

```
PyTorch version: 1.13.0.dev20220919                                                                      
Is debug build: False                                                                                    
CUDA used to build PyTorch: 11.6                                                                         
ROCM used to build PyTorch: N/A                                                                          
                                                                                                         
OS: Ubuntu 20.04.4 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.16.3
Libc version: glibc-2.31

Python version: 3.9.13 (main, Aug 25 2022, 23:26:10)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-125-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.6.124
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000

Nvidia driver version: 510.68.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True 

Versions of relevant libraries:
[pip3] functorch==0.2.1
[pip3] numpy==1.23.1
[pip3] torch==1.13.0.dev20220919
[pip3] torchaudio==0.13.0.dev20220919
[pip3] torchvision==0.14.0.dev20220919
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] functorch                 0.2.1                    pypi_0    pypi
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0 
[conda] numpy                     1.23.1           py39h6c91a56_0  
[conda] numpy-base                1.23.1           py39ha15fc14_0  
[conda] pytorch                   1.13.0.dev20220919 py3.9_cuda11.6_cudnn8.3.2_0    pytorch-nightly
[conda] pytorch-cuda              11.6                 h867d48c_0    pytorch-nightly
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.13.0.dev20220916+cpu          pypi_0    pypi
[conda] torchaudio                0.13.0.dev20220919      py39_cu116    pytorch-nightly
[conda] torchvision               0.14.0.dev20220916+cpu          pypi_0    pypi
```