1. 先检查你的CUDA环境
检查NVIDIA驱动和CUDA版本
nvidia-smi

查看右上角的CUDA Version(如:12.8)
2. 根据CUDA版本安装对应的PyTorch
访问 PyTorch官网 获取正确的安装命令。https://pytorch.org/get-started/locally/

pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu128
卸载原有 torch torchvision ,执行以上命令安装。

3. 验证安装
import torch
检查PyTorch版本
print(torch.version)
检查CUDA是否可用
print(torch.cuda.is_available()) # 应该返回True
检查CUDA版本
print(torch.version.cuda)
检查GPU数量
print(torch.cuda.device_count())
执行 python 代码训练模型,显示 NVIDIA GPU 运行应用程序
