How to install CUDA enabled PyTorch in a Docker container?

I got it working after many, many tries. Posting the answer here in case it helps anyone.

Basically, I installed pytorch and torchvision through pip (from within the conda environment) and rest of the dependencies through conda as usual.

This is how the final Dockerfile looks:

# Use nvidia/cuda imageFROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04# set bash as current shellRUN chsh -s /bin/bashSHELL ["/bin/bash", "-c"]# install anacondaRUN apt-get updateRUN apt-get install -y wget bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 git mercurial subversion && \        apt-get cleanRUN wget --quiet https://repo.anaconda.com/archive/Anaconda3-2020.02-Linux-x86_64.sh -O ~/anaconda.sh && \        /bin/bash ~/anaconda.sh -b -p /opt/conda && \        rm ~/anaconda.sh && \        ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \        echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \        find /opt/conda/ -follow -type f -name '*.a' -delete && \        find /opt/conda/ -follow -type f -name '*.js.map' -delete && \        /opt/conda/bin/conda clean -afy# set path to condaENV PATH /opt/conda/bin:$PATH# setup conda virtual environmentCOPY ./requirements.yaml /tmp/requirements.yamlRUN conda update conda \    && conda env create --name camera-seg -f /tmp/requirements.yamlRUN echo "conda activate camera-seg" >> ~/.bashrcENV PATH /opt/conda/envs/camera-seg/bin:$PATHENV CONDA_DEFAULT_ENV $camera-seg

And this is how the requirements.yaml looks like:

name: camera-segchannels:  - defaults  - conda-forgedependencies:  - python=3.6  - pip  - numpy  - pillow  - yaml  - pyyaml  - matplotlib  - jupyter  - notebook  - tensorboardx  - tensorboard  - protobuf  - tqdm  - pip:    - torch    - torchvision

Then I build the container using the command docker build -t camera-seg . and PyTorch is now being able to recognize CUDA.