Set Pytorch to run on AMD GPU

Question

According to the official docs, now PyTorch supports AMD GPUs. ROCm 4.2 can be installed through pip. But I can not find in Google nor the official docs how to force my DL training to use the GPU. What is the AMD equivalent to the following command?

torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Answer 1

ROCm translates CUDA codes. So torch.cuda.is_available() should work.