想必大家查看自己cuda是11.2版本,但是在pytorch官网没有找到对应的安装命令,请参考一下命令:

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

测试:

import torch

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

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

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

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

Logo

欢迎来到FlagOS开发社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。

更多推荐