Update Dockerfile

Dockerfile

@@ -21,7 +21,8 @@ WORKDIR /code
 COPY ./requirements.txt /code/requirements.txt
 
 # make sure cudnn errors go away
-RUN export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:`python3 -c 'import os; import torch; print(os.path.dirname(torch.__file__) +"/lib")'`
+#RUN export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:`python3 -c 'import os; import torch; print(os.path.dirname(torch.__file__) +"/lib")'`
+RUN export LD_LIBRARY_PATH=`python3 -c 'import os; import nvidia.cublas.lib; import nvidia.cudnn.lib; import torch; print(os.path.dirname(nvidia.cublas.lib.__file__) + ":" + os.path.dirname(nvidia.cudnn.lib.__file__) + ":" + os.path.dirname(torch.__file__) +"/lib")'`:$LD_LIBRARY_PATH
 
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user