Repository: pytorch
Issue: Training fails with Torch 2.1.0 on Nvidia Jetpack 5.1.2
Labels: triaged,module: jetson

### 🐛 Describe the bug

Hello

We are trying to run a training on Nvidia Jetson devices with compute capabilities 7.2 and 8.7.
The system properties are as follows:
```
Python 3.8
Torch 2.1.0
Torchvision 0.16.2
CUDA 11.4
Nvidia Jetpack 5.1.2
Ubuntu 20.04
```

At the begining of a simple MNIST training, while executing `loss.backward()` we get the error below:
```
File "/mnt/nvme/.venvs/venv3_8/lib/python3.8/site-packages/torch/autograd/__init__.py", line 204, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Event device type CUDA does not match blocking stream's device type CPU.
```
The error occurs when we use an environment in which Torch is added using the [.whl from Jetson Zoo](https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048).
Eventhough we build our own torch .whl according to the [build file in Jetson containers](https://github.com/dusty-nv/jetson-containers/blob/master/packages/pytorch/build.sh) we get the same error.

When we use the same scripts to build torch for CUDA 12.2 and run the same simple MNIST training we do not get the error.

I appreciate any help.

### Versions

Collecting environment information...
PyTorch version: 2.1.0a0+git7bcf7da
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (aarch64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.31.2
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov  7 2024, 13:10:47)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.10.120-tegra-aarch64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.4.315
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/aarch64-linux-gnu/libcudnn.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_train.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_train.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    aarch64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
CPU(s):                          12
On-line CPU(s) list:             0-11
Thread(s) per core:              1
Core(s) per socket:              4
Socket(s):                       3
Vendor ID:                       ARM
Model:                           1
Model name:                      ARMv8 Processor rev 1 (v8l)
Stepping:                        r0p1
CPU max MHz:                     2201.6001
CPU min MHz:                     115.2000
BogoMIPS:                        62.50
L1d cache:                       768 KiB
L1i cache:                       768 KiB
L2 cache:                        3 MiB
L3 cache:                        6 MiB
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; CSV2, but not BHB
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm lrcpc dcpop asimddp uscat ilrcpc flagm

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] onnx==1.17.0
[pip3] torch==2.1.0a0+git7bcf7da
[pip3] torchaudio==2.1.0+6ea1133
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.16.2+c6f3977
[conda] Could not collect

cc @ptrblck @puririshi98