diff --git a/training/test_train.py b/training/test_train.py
index b054cfa..6d64ea2 100644
--- a/training/test_train.py
+++ b/training/test_train.py
@@ -35,14 +35,18 @@
 def model(request):
     # Load dataset
     data = None
+    prefix = ""
     try:
         data = pd.read_csv(request.param["dataset"])
     except FileNotFoundError:
         # Check if the dataset is in the parent directory
-        data = pd.read_csv("../" + request.param["dataset"])
+        prefix = "../"
+        data = pd.read_csv(prefix + request.param["dataset"])
 
     # Load TF model from SavedModel
-    sqli_model = TFSMLayer(request.param["model_path"], call_endpoint="serving_default")
+    sqli_model = TFSMLayer(
+        prefix + request.param["model_path"], call_endpoint="serving_default"
+    )
 
     # Tokenize the sample
     tokenizer = Tokenizer(num_words=MAX_WORDS, filters="")