Thanks for following up. I’m sorry for the late reply. This is the info about my environment.
I’m in the process of creating a new environment with the latest Pytorch
Collecting environment information...
PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.5 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.1.168
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti
GPU 7: GeForce RTX 2080 Ti

Nvidia driver version: 460.80
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.5.0

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] pytorch-lightning==0.9.0
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas                      2.21                        mkl    conda-forge
[conda] cudatoolkit               10.1.243             h036e899_6    conda-forge
[conda] libblas                   3.8.0                    21_mkl    conda-forge
[conda] libcblas                  3.8.0                    21_mkl    conda-forge
[conda] liblapack                 3.8.0                    21_mkl    conda-forge
[conda] liblapacke                3.8.0                    21_mkl    conda-forge
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.19.1           py38hbc27379_2    conda-forge
[conda] pytorch                   1.6.0           py3.8_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] pytorch-lightning         0.9.0                      py_0    conda-forge
[conda] torchvision               0.7.0                py38_cu101    pytorch