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

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