import torch

def test_cuda_availability():
    """全面测试CUDA可用性"""
    
    print("="*50)
    print("PyTorch CUDA 测试")
    print("="*50)
    
    # 基本信息
    print(f"PyTorch版本: {torch.__version__}")
    print(f"CUDA可用: {torch.cuda.is_available()}")
    
    if not torch.cuda.is_available():
        print("CUDA不可用，可能原因：")
        print("1. 未安装CUDA驱动")
        print("2. 安装的是CPU版本的PyTorch")
        print("3. CUDA版本与PyTorch不匹配")
        return False
    
    # 设备信息
    device_count = torch.cuda.device_count()
    print(f"发现 {device_count} 个CUDA设备")
    
    for i in range(device_count):
        print(f"\n设备 {i}:")
        print(f"  名称: {torch.cuda.get_device_name(i)}")
        print(f"  内存总量: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")
        print(f"  计算能力: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}")
    
    # 简单张量测试
    print("\n运行CUDA测试...")
    try:
        x = torch.randn(3, 3).cuda()
        y = torch.randn(3, 3).cuda()
        z = x + y
        print("CUDA计算测试: 成功!")
        print(f"设备上的张量形状: {z.shape}")
        return True
    except Exception as e:
        print(f"CUDA计算测试失败: {e}")
        return False

if __name__ == "__main__":
    test_cuda_availability()