diff --git a/docker/build-torchbench-nightly-docker.sh b/docker/build-torchbench-nightly-docker.sh
index 5194affc4c..8f53442dd9 100644
--- a/docker/build-torchbench-nightly-docker.sh
+++ b/docker/build-torchbench-nightly-docker.sh
@@ -1 +1,4 @@
-docker build . -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:latest
+TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH:-main}
+
+docker build . -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:latest \
+    --build-arg TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH}
diff --git a/docker/torchbench-nightly.dockerfile b/docker/torchbench-nightly.dockerfile
index c48436f47f..a35a3dda20 100644
--- a/docker/torchbench-nightly.dockerfile
+++ b/docker/torchbench-nightly.dockerfile
@@ -8,10 +8,13 @@ ENV SETUP_SCRIPT=/workspace/setup_instance.sh
 ARG TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH:-main}
 ARG FORCE_DATE=${FORCE_DATE}
 
-# Setup Conda env and CUDA
+# Checkout Torchbench and submodules
 RUN git clone -b "${TORCHBENCH_BRANCH}" --single-branch \
- https://github.com/pytorch/benchmark /workspace/benchmark
+    https://github.com/pytorch/benchmark /workspace/benchmark
+RUN cd /workspace/benchmark \
+    git submodule update --init --recursive
 
+# Setup conda env and CUDA
 RUN cd /workspace/benchmark && \
     . ${SETUP_SCRIPT} && \
     python ./utils/python_utils.py --create-conda-env ${CONDA_ENV} && \
diff --git a/userbenchmark/triton/install.py b/userbenchmark/triton/install.py
index 99051367e2..360806c3d2 100644
--- a/userbenchmark/triton/install.py
+++ b/userbenchmark/triton/install.py
@@ -9,7 +9,7 @@
 
 def install_fbgemm():
     cmd = [sys.executable, "setup.py", "bdist_wheel", "--package_variant=genai"]
-    subprocess.check_call(cmd, cwd=FBGEMM_PATH)
+    subprocess.check_call(cmd, cwd=str(FBGEMM_PATH.resolve()))
 
 def test_fbgemm():
     cmd = [sys.executable, "-c", '"import fbgemm_gpu.experimental.gen_ai"']