Ubuntu的pip安装pytorch-gpu版

1. 显卡驱动要装好，确保nvidia-smi可以运行

新的gpu驱动已经可以支持cuda12.4

2. 推荐在conda环境中用pip进行安装，conda环境配置可以看下述博客

Ubuntu安装miniconda3_tiantangzhixia的博客-CSDN博客Ubuntu安装miniconda3_ubuntu安装miniconda3https://blog.csdn.net/baidu_35280192/article/details/127530088?csdn_share_tail=%7B%22type%22%3A%22blog%22%2C%22rType%22%3A%22article%22%2C%22rId%22%3A%22127530088%22%2C%22source%22%3A%22baidu_35280192%22%7D

3. 进入到conda虚拟环境中，安装torch环境

截至到目前的最新版是torch2.6.0，可以使用cuda12.4。推荐采用python3.13。

conda create -n torch python=3.13 -y
conda activate torch
pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
pip install torchmetrics tensorboard tensorboardX netcdf4 h5py haversine scipy tqdm pandas pyyaml

4. 检查是否可以调用gpu

python
import torch
torch.cuda.is_available()

如果是True，那就ok啦