Check if CUDA works properly:

import torch 

torch.cuda.is_available()

Or simply create a tensor on the GPU:

tensor = torch.tensor([1.0, 2.0]).cuda()