The instructions to install Cuda look right.
Make sure to pip uninstall torch before reinstalling it.
But here again, there is 99% chance that this is a bug in your code (or the inputs/dataset) and not an installation problem.