Jetpack 5.1.2 Model Training Fails with Torch 2.1.0

Robotics & Edge Computing Jetson & Embedded Systems Jetson AGX Xavier
8

Hi,

According to this:

RuntimeError: Event device type CUDA does not match blocking stream's device type CPU.

It might not relate to the compatibility between PyTorch and TorchVision.
Could you share how do you apply the training?

Please try to add the below line into your script to see if it helps.

import pycuda.autoinit

Thanks.