FROM jupyter/datascience-notebook:65761486d5d3

RUN conda install -y pytorch=1.0.1 torchvision=0.2.2 cudatoolkit=9.0 -c pytorch

RUN pip install \
    torchsummary==1.5.1 \
    visdom==0.1.8.8

ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
ENV NVIDIA_REQUIRE_CUDA "cuda>=9.0"
