@heerah1493 you do not need to install CUDA 11.8 to enable GPU acceleration in TensorFlow/PyTorch, and in fact at first you should probably stick with the version of CUDA that came with JetPack because that is what those pip wheels were built against.
Seeing as you now are also getting cuDNN failures, and are also on one of the earlier versions of JetPack 5, I would recommend re-flashing your board with the latest JetPack (at this time, that’s JetPack 5.1.1) and starting with a fresh environment. Then install PyTorch like in this thread:
You can also try using the l4t-pytorch, l4t-tensorflow, or l4t-ml containers which already come with PyTorch/TensorFlow pre-installed. You should be able to do this without reflashing, as CUDA/cuDNN/ect are installed inside those containers.
Hi @heerah1493, I would recommend posting a new topic about this specific issue along with references to the code that you are running, or to file an issue with the upstream maintainers of that YOLOv5 project.