search

Googolxx / STF

Pytorch implementation of the paper "The Devil Is in the Details: Window-based Attention for Image Compression".
Apache License 2.0
161 stars 20 forks source link

Using A100 GPU to train the model #1

Closed wangfeng9512 closed 2 years ago

wangfeng9512 commented 2 years ago

When using A100 GPU to train the model "cnn", maybe we will get the following error:

NVIDIA A100-PCIE-40GB with CUDA capability sm_80 is not compatible with the current PyTorch installation. The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_70. If you want to use the NVIDIA A100-PCIE-40GB GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

wangfeng9512 commented 2 years ago

I find the following command can solve the problem:

pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

We can use the above command after ”pip install -e '.[dev]'“.

Reference link: https://discuss.pytorch.org/t/nvidia-a100-gpu-runtimeerror-cudnn-error-cudnn-status-mapping-error/121648