Skip to content

Commit

Permalink
Use official pytorch base image
Browse files Browse the repository at this point in the history
  • Loading branch information
Scott Davidson committed Sep 22, 2023
1 parent 4fd0031 commit 802f1d8
Showing 1 changed file with 10 additions and 9 deletions.
19 changes: 10 additions & 9 deletions images/pytorch-benchmarks/Dockerfile
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
FROM nvcr.io/nvidia/pytorch:22.12-py3
FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

# RUN pip install -U pip
# RUN pip install torchvision torchaudio
RUN pip uninstall torch -y && pip install --no-cache torch torchvision torchaudio
RUN apt update && apt install -y git
# TODO: Checkout certain commit here
RUN git clone https://github.com/pytorch/benchmark
WORKDIR /workspace/benchmark
# Add other benchmarks here?
RUN python install.py alexnet resnet50 llama

# RUN git clone https://github.com/pytorch/benchmark
# WORKDIR /workspace/benchmark
# # Add other benchmarks here?
# RUN python install.py alexnet resnet50 llama
# PyTorch install.py pins numpy=1.21.2 but this breaks numba so update both here
RUN pip install -U numpy numba

# COPY run-benchmark.sh /usr/local/bin/
COPY run-benchmark.sh /usr/local/bin/

0 comments on commit 802f1d8

Please sign in to comment.