You recommand installing pytorch with pip.
But I found that the default installation is the cpu version of PyTorch.
I tried to test your code and got the error:
AssertionError: Torch not compiled with CUDA enabled
It is recommended to specify the GPU version of PyTorch to be used in the documentation
BTW I used this code to install PyTorch and fix the cuda error reporting problem
pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio===0.8.0 -f https://download.pytorch.org/whl/torch_stable.html
You recommand installing pytorch with pip. But I found that the default installation is the cpu version of PyTorch.
I tried to test your code and got the error:
AssertionError: Torch not compiled with CUDA enabled
It is recommended to specify the GPU version of PyTorch to be used in the documentation BTW I used this code to install PyTorch and fix the cuda error reporting problem
pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio===0.8.0 -f https://download.pytorch.org/whl/torch_stable.html