diff --git a/hpvm/test/epoch_dnn/main.py b/hpvm/test/epoch_dnn/main.py
index fe4408ef6912610066872e68192d796d997e301f..e94b13f48ce1451cca218980ff63eb1ed392ff1f 100644
--- a/hpvm/test/epoch_dnn/main.py
+++ b/hpvm/test/epoch_dnn/main.py
@@ -70,6 +70,9 @@ SCP_HOST = "root@espgate.cs.columbia.edu"
 SCP_PWD = "openesp"
 SCP_DST = "~/NV_NVDLA"
 
+if len(sys.argv) > 1:
+    WORKING_DIR = Path(sys.argv[1])
+
 # Reproducibility
 np.random.seed(42)