From 78307b0e8beb5d41ee003be8b5db168c2b3ef4e2 Mon Sep 17 00:00:00 2001
From: Gabriele Oliaro <goliaro@cs.cmu.edu>
Date: Tue, 26 Nov 2024 19:13:07 +0000
Subject: [PATCH] update

---
 docker/run.sh | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/docker/run.sh b/docker/run.sh
index 759da521aa..62d7468a00 100755
--- a/docker/run.sh
+++ b/docker/run.sh
@@ -127,8 +127,7 @@ fi
 
 ssh_key_volume=""
 ssh_key_path="$HOME/.ssh/id_rsa"
-if [ -f "$ssh_key_path" ]; then
-  # If the token exists, add the volume mount to the Docker command
-  ssh_key_volume+="-v $ssh_key_path:/root/.ssh/id_rsa"
+if [ -f "$ssh_key_path" ] && [ -f "$ssh_key_path.pub" ]; then
+  ssh_key_volume="-v $ssh_key_path:/root/.ssh/id_rsa -v $ssh_key_path.pub:/root/.ssh/id_rsa.pub"
 fi
 eval docker run -it "$gpu_arg" "--shm-size=${SHM_SIZE}" "--cap-add=SYS_PTRACE" "${ssh_key_volume}" "${hf_token_volume}" "${port_forward_arg}" "${image}-${FF_GPU_BACKEND}${gpu_backend_version}:latest"