Thanks for your advice. I also test my code on 1.10.1. The print("Is cuDNN version:", torch.backends.cudnn.version()) gives me 8200. It still gives me nan error. The collected environment is below.
Collecting environment information...
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.6 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: Could not collect
CMake version: version 3.19.6
Libc version: glibc-2.10

Python version: 3.7.4 (default, Aug 13 2019, 20:35:49)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-4.4.0-142-generic-x86_64-with-debian-stretch-sid
Is CUDA available: True
CUDA runtime version: Could not collect

GPU models and configuration:
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000

Nvidia driver version: 465.19.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.10.1
[pip3] torchaudio==0.10.1+cu113
[pip3] torchvision==0.11.2
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] magma-cuda113             2.5.2                         1    pytorch
[conda] mkl                       2021.3.0           h06a4308_520
[conda] mkl-include               2021.3.0           h06a4308_520
[conda] mkl-service               2.4.0            py37h7f8727e_0
[conda] mkl_fft                   1.3.1            py37hd3c417c_0
[conda] mkl_random                1.2.2            py37h51133e4_0
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] pytorch                   1.10.1          py3.7_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.10.1+cu113             pypi_0    pypi
[conda] torchvision               0.11.2               py37_cu113    pytorch