File size: 249 Bytes
b9cd341
8d7f55f
554e451
 
 
 
 
 
1
2
3
4
5
6
7
8
9
import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
else:
    device = torch.device("cpu")
    print("CUDA is not available, using CPU.")