FROM pytorch/pytorch:1.6.0-cuda10.1-cudnn7-devel
# Need rsync for ray
RUN apt-get update && \
    apt-get -y install rsync && \
    rm -rf /var/lib/apt/lists/*
RUN pip install pytorch-lightning==0.9.0 pytorch-lightning-bolts==0.2.2 ray[tune]==0.8.7 hydra-core==1.0.3 munch wandb==0.9.6 && \
    rm -rf /root/.cache/pip
