I am trying to install torch package and the verification code runs successfully as follows:

I am trying to run the model, the output is the runtime error ("all CUDA-capable devices are busy or unavailable")
My setup is the following:
Ubuntu 18.04.6 LTS (GNU/Linux 5.4.0-122-generic x86_64)

- NVIDIA A100 40G Pcle *2
- CUDA : 11.0
- cuDNN : 8.4.1
- python : 3.9
- pytorch : pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f - https://download.pytorch.org/whl/torch_stable.html
I tried many solutions from the stack overflow, but it didn't work. I don't want to reboot every time it happens.
This is an example of code that may fail.

I don't know what is wrong. How does one use Pytorch (+ cuda) with an A100 GPU? If you have experience installing in a similar development environment, please give me some advice....