RTX3080TiがPyTorchから認識されなくて困った件

起こった問題

RTX3080TiがPyTorchの最新バージョンの1.10.1から認識されなくて困った。CUDAは11以降を使っている。

動作環境

• Ubuntu20.04上のdockerコンテナ内で作業
• dockerイメージ: nvidia/cuda:11.3.0-devel-ubuntu20.04
• cuda: 11.3
• cuDNN: 上記cudaに対応したもの
• Nvidia driver: 495.46

Webで何件が見つけたこのエラーが出てしまう。

GeForce RTX 3080Ti with CUDA capability sm_86 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_70 sm_75.
If you want to use the GeForce RTX 3080Ti GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

PYTorchから対応しているバージョンを出してみてもsm_70までしか利用できないようだ。RTX3080Tiを使うにはsm_86が必要。

import torch
print(torch.cuda.get_arch_list())
['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75']

解決方法

いろいろ試した結果、PyTorchのPreview版を使うことで解決

こちらのサイトでPreviewを選ぶとStableよりも新しいバージョンのインストールコマンドが出てくるので、そちらを利用した（表示されたエラーメッセージは正しかった）
https://pytorch.org/get-started/locally/

動いている環境での、インストールされているPyTorchのバージョンは以下のようになっていた。

# pip list | grep torch
torch                              1.11.0.dev20220124+cu113
torchaudio                         0.11.0.dev20220124+cu113
torchvision                        0.12.0.dev20220124+cu113

torch.cuda.get_arch_list()の結果も変わっていた。

import torch
print(torch.cuda.get_arch_list())
['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86']