FROM nvcr.io/nvidia/pytorch:22.11-py3

RUN pip install torch-pruning==1.3.5

ADD . /workspace

WORKDIR /workspace

RUN wget https://github.com/VainF/Torch-Pruning/releases/download/v1.1.4/cifar10_resnet56.pth
RUN wget https://github.com/VainF/Torch-Pruning/releases/download/v1.1.4/cifar100_vgg19.pth

CMD ["sh","run_all.sh"]