diff --git a/training.py b/training.py
index 1899ef2..b4ca95e 100755
--- a/training.py
+++ b/training.py
@@ -6,30 +6,39 @@
 
 import numpy as np
 import pandas as pd
+from sklearn.metrics import accuracy_score
 import torch
 from torch import nn, optim
 import torch.nn.functional as F
 import sys
+from sacred import Experiment
+from sacred.observers import FileStorageObserver
 
 
 # In[ ]:
 
 
-epochs = int(sys.argv[1])
+ex = Experiment("file_observer", save_git_info=False)
+ex.observers.append(FileStorageObserver('my_runs'))
+
+@ex.config
+def my_config():
+    epochs = 400
 
 
 # In[ ]:
 
 
-X_train = pd.read_csv('X_train.csv')
-y_train = pd.read_csv('y_train.csv')
-
-
-# In[ ]:
-
-
-X_train = torch.from_numpy(np.array(X_train)).float()
-y_train = torch.squeeze(torch.from_numpy(y_train.values).float())
+def prepare_data():
+    X_train = pd.read_csv('X_train.csv')
+    y_train = pd.read_csv('y_train.csv')
+    X_train = torch.from_numpy(np.array(X_train)).float()
+    y_train = torch.squeeze(torch.from_numpy(y_train.values).float())
+    X_test = pd.read_csv('X_test.csv')
+    y_test = pd.read_csv('y_test.csv')
+    X_test = torch.from_numpy(np.array(X_test)).float()
+    y_test = torch.squeeze(torch.from_numpy(y_test.values).float())
+    return X_train, y_train
 
 
 # In[ ]:
@@ -50,51 +59,75 @@ class Net(nn.Module):
 # In[ ]:
 
 
-net = Net(X_train.shape[1])
-criterion = nn.BCELoss()
-optimizer = optim.Adam(net.parameters(), lr=0.001)
-
-
-# In[ ]:
-
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-X_train = X_train.to(device)
-y_train = y_train.to(device)
-
-net = net.to(device)
-criterion = criterion.to(device)
-
-
-# In[ ]:
-
-
 def calculate_accuracy(y_true, y_pred):
     predicted = y_pred.ge(.5).view(-1)
     return (y_true == predicted).sum().float() / len(y_true)
 
+
+# In[ ]:
+
+
 def round_tensor(t, decimal_places=3):
     return round(t.item(), decimal_places)
 
+# In[ ]:
 
-for epoch in range(epochs):
-    y_pred = net(X_train)
-    y_pred = torch.squeeze(y_pred)
-    train_loss = criterion(y_pred, y_train)
-    if epoch % 100 == 0:
-        train_acc = calculate_accuracy(y_train, y_pred)
-        print(
-            f'''epoch {epoch}
-            Train set - loss: {round_tensor(train_loss)}, accuracy: {round_tensor(train_acc)}
-            ''')
+
+def train_model(X_train, y_train, device, epochs):
+    net = Net(X_train.shape[1])
+    criterion = nn.BCELoss()
+    optimizer = optim.Adam(net.parameters(), lr=0.001)
+    
+    X_train = X_train.to(device)
+    y_train = y_train.to(device)
+    
+    net = net.to(device)
+    criterion = criterion.to(device)
+
+    for epoch in range(epochs):
+        y_pred = net(X_train)
+        y_pred = torch.squeeze(y_pred)
+        train_loss = criterion(y_pred, y_train)
+        if epoch % 100 == 0:
+            train_acc = calculate_accuracy(y_train, y_pred)
+            print(
+                f'''epoch {epoch}
+                Train set - loss: {round_tensor(train_loss)}, accuracy: {round_tensor(train_acc)}
+                ''')
     optimizer.zero_grad()
     train_loss.backward()
     optimizer.step()
+    return net
 
 
 # In[ ]:
 
 
-torch.save(net, 'model.pth')
+def evaluate_model(X_test, y_test, device):
+    X_test = X_test.to(device)
+    y_test = y_test.to(device)
+    net = torch.load('model.pth')
+    y_pred = net(X_test)
+    y_pred = y_pred.ge(.5).view(-1).cpu()
+    y_test = y_test.cpu()
+    accuracy = accuracy_score(y_test, y_pred)
+    return accuracy
+
+
+# In[ ]:
+
+
+@ex.automain
+def my_main(epochs, _run):
+    _run.info["epochs"] = epochs
+    X_train, X_test, y_train, y_test = prepare_data()
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    
+    model = train_model(X_train, y_train, device, epochs)
+    torch.save(model, 'model.pth')
+    ex.add_artifact('model.pth')
+    
+    accuracy = evaluate_model(X_test, y_test, device)
+    print(accuracy)
+    _run.info["accuracy"] = accuracy