import torch

print("torch.__version__: ", torch.__version__)
print("torch.version.cuda: ", torch.version.cuda)
print("torch.cuda.is_available: ", torch.cuda.is_available())
print("torch.backends.cudnn.version: ", torch.backends.cudnn.version())

if torch.cuda.is_available():
    device_name = torch.cuda.get_device_name(0)
    device_count = torch.cuda.device_count()
    print('torch.cuda.get_device_name: ' + device_name)
    print('torch.cuda.device_count: ' + str(device_count))

    for i in range(device_count):
        props = torch.cuda.get_device_properties(i)
        print(f"\n设备 {i}:")
        print(f"  名称: {props.name}")
        print(f"  计算能力: {props.major}.{props.minor}")
        print(f"  总内存: {props.total_memory / (1024 ** 2):.2f} MB")
        print(f"  多处理器数量: {props.multi_processor_count}")
        print(f"  每个多处理器的最大线程数: {props.max_threads_per_multi_processor}")
        print(f"  是否集成在主板上: {props.is_integrated}")

    # 测试深度学习模型是否可以在GPU上训练
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"\n当前设备: {device}")

    # 创建一个示例张量并移动到当前设备上
    tensor = torch.randn(3, 3).to(device)
    print(f"张量在{tensor.device}上")

    # 如果张量在GPU上,则可以在GPU上进行训练
    if tensor.device.type == "cuda":
        print("深度学习模型可以在GPU上进行训练")
    else:
        print("深度学习模型无法在GPU上进行训练,将使用CPU进行训练")
else:
    print("此机器上没有 CUDA 设备可用")