diff --git a/hpvm/test/epoch_dnn/main.py b/hpvm/test/epoch_dnn/main.py
index e94b13f48ce1451cca218980ff63eb1ed392ff1f..2ecfa5343aa65c3db80c1b561379f7a11a8c4e69 100644
--- a/hpvm/test/epoch_dnn/main.py
+++ b/hpvm/test/epoch_dnn/main.py
@@ -2,6 +2,7 @@ import site
 from pathlib import Path
 from os import makedirs
 import numpy as np
+import sys
 
 import torch
 from torch2hpvm import BinDataset, ModelExporter