cuda11.2 corresponde a la instalación de pytorch

Debe verificar que su versión de cuda sea 11.2, pero no puede encontrar el comando de instalación correspondiente en el sitio web oficial de pytorch. Consulte el comando:

conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.2 -c pytorch -c conda-forge

prueba:

import torch

torch.cuda.is_available()       # cuda是否可用

torch.cuda.current_device()     # 返回当前设备索引

torch.cuda.device_count()       # 返回GPU的数量

torch.cuda.get_device_name(0)   # 返回gpu名字，设备索引默认从0开始