import torch
# 检查CUDA是否可用
cuda_available = torch.cuda.is_available()
if cuda_available:
# 获取GPU设备数量
num_gpu = torch.cuda.device_count()
# 获取当前使用的GPU索引
current_gpu_index = torch.cuda.current_device()
# 获取当前GPU的名称
current_gpu_name = torch.cuda.get_device_name(current_gpu_index