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.")