diff --git a/src/scripts/paths.py b/src/scripts/paths.py
index 8c9434e..939c710 100644
--- a/src/scripts/paths.py
+++ b/src/scripts/paths.py
@@ -25,5 +25,6 @@
 # Absolute path to the `src/tex/figures` folder (contains figure output)
 figures = tex / "figures"
 
-# Absolute path to the `src/tex/output` folder (contains other user-defined output)
-output = tex / "output"
\ No newline at end of file
+# Absolute path to the `src/tex/output` folder
+# (contains other user-defined output)
+output = tex / "output"
diff --git a/src/scripts/train.py b/src/scripts/train.py
index a24465b..3bba349 100644
--- a/src/scripts/train.py
+++ b/src/scripts/train.py
@@ -1,7 +1,9 @@
 """
 Simple stubs to use for re-train of the final model
-Can leave a default data source, or specify that 'load data' loads the dataset used in the final version
+Can leave a default data source, or specify that 'load data' loads the dataset
+used in the final version
 """
+
 import argparse
 
 
@@ -11,12 +13,14 @@ def architecture():
     """
     return 0
 
+
 def load_data(data_source):
     """
     :return: data loader or full training data, split in val and train
     """
     return 0, 0
 
+
 def train_model(data_source, n_epochs):
     """
     :param data_source:
@@ -26,13 +30,17 @@ def train_model(data_source, n_epochs):
     data = load_data(data_source)
     model = architecture()
 
-    return 0
+    return data, model
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--data_source", type=str, help="Data used to train the model")
-    parser.add_argument("--n_epochs", type=int, help='Integer number of epochs to train the model')
+    parser.add_argument("--data_source", type=str,
+                        help="Data used to train the model")
+    parser.add_argument(
+        "--n_epochs", type=int,
+        help="Integer number of epochs to train the model"
+    )
 
     args = parser.parse_args()