diff --git a/images/pytorch-benchmarks/Dockerfile b/images/pytorch-benchmarks/Dockerfile
index 50ccab7..d7d79ca 100644
--- a/images/pytorch-benchmarks/Dockerfile
+++ b/images/pytorch-benchmarks/Dockerfile
@@ -1,6 +1,6 @@
 FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
 
-RUN apt update && apt install -y git
+RUN apt update && apt install -y git time
 # TODO: Checkout certain commit here
 RUN git clone https://github.com/pytorch/benchmark
 WORKDIR /workspace/benchmark
diff --git a/images/pytorch-benchmarks/run-benchmark.sh b/images/pytorch-benchmarks/run-benchmark.sh
deleted file mode 100755
index 11cfcaa..0000000
--- a/images/pytorch-benchmarks/run-benchmark.sh
+++ /dev/null
@@ -1,3 +0,0 @@
-#!/bin/bash
-# Env vars should be set in container spec
-python run.py $BENCHMARK_NAME -t $BENCHMARK_TYPE -d $BENCHMARK_DEVICE --bs $INPUT_DATA_BATCH_SIZE
\ No newline at end of file