import torch

# 获取 PyTorch 版本
pytorch_version = torch.__version__
# 获取 CUDA 版本
cuda_version = torch.version.cuda
# 获取 cuDNN 版本
cudnn_version = torch.backends.cudnn.version()

print(f"PyTorch 版本: {pytorch_version}")
print(f"CUDA 版本: {cuda_version}")
print(f"cuDNN 版本: {cudnn_version}")

# 检查是否可以访问 CUDA
if torch.cuda.is_available():
    print("pip install sentence-transformers -i https://pypi.mirrors.ustc.edu.cn/simpleCUDA is available. GPU name:", torch.cuda.get_device_name(0))
else:
    print("CUDA is not available. Please check your installation.")