18 lines
606 B
Python
18 lines
606 B
Python
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.") |