I have reduced the batch size but it doesn’t work and here is the code snippet
OS: Ubuntu 20.04.2 LTS
GCC version: (Ubuntu 10.2.0-5ubuntu1~20.04) 10.2.0
CMake version: version 3.16.3
Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1080 Ti
Nvidia driver version: 465.19.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.8.0.4
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.5.0
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.6.0a0+82fd1c8
[conda] blas                      1.0                         mkl    defaults
[conda] mkl                       2020.2                      256    defaults
[conda] mkl-service               2.3.0            py36he8ac12f_0    defaults
[conda] mkl_fft                   1.3.0            py36h54f3939_0    defaults
[conda] mkl_random                1.1.1            py36h0573a6f_0    defaults
[conda] pytorch                   1.5.0           py3.6_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] torchsummary              1.5.1                    pypi_0    pypi
[conda] torchvision               0.6.0                py36_cu101    pytorch