PyTorchでRTX 3090を使う
※2020/10/24日時点の情報です。
10/28日にPyTorch 1.7が公開され，それに伴いCUDA11に対応・RTX3090でも利用可能となっています。
そのため，単にPyTorch 1.7に上げればRTX3080/3090でも使えます。
apexのインストールは，以下で出来ました。
https://github.com/NVIDIA/apex/issues/988#issuecomment-726343453
以下は古い情報です。
課題
PyTorch 1.6では，現状RTX 3080/3090に対応しておらず，以下のようなエラーを吐きます。
GeForce RTX 3090 with CUDA capability sm_86 is not compatible with the current PyTorch installation.
解決方法
対処方法は，開発版のPyTorch 1.8を入れることです。
pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu110/torch_nightly.html
cu110となっていますが，cuda11.1を使っていても動くので問題ありません。
バージョン情報
NVIDIA-SMI 455.23.05 Driver Version: 455.23.05 CUDA Version: 11.1/Ubuntu 20.04
結果
以前は2070 SUPERを使っていたのですが，学習速度が3倍程度早くなりました。あと，jaxは普通に動きました。
Ref.
Author And Source
この問題について(PyTorchでRTX 3090を使う), 我々は、より多くの情報をここで見つけました https://qiita.com/atksh/items/10e99066cbefc0876704著者帰属：元の著者の情報は、元のURLに含まれています。著作権は原作者に属する。
Content is automatically searched and collected through network algorithms . If there is a violation . Please contact us . We will adjust (correct author information ,or delete content ) as soon as possible .