diff --git a/Biblioteka_DL/dllib-mlflow.py b/Biblioteka_DL/dllib-mlflow.py
index 36c3bde..a00049e 100644
--- a/Biblioteka_DL/dllib-mlflow.py
+++ b/Biblioteka_DL/dllib-mlflow.py
@@ -291,7 +291,8 @@ def my_main(epochs):
     loss_fn = nn.CrossEntropyLoss()
     # epochs = 1000
     # epochs = epochs
-    
+    mlflow.log_param("epochs", epochs)
+
     def print_(loss):
         print ("The loss calculated: ", loss)
 
@@ -304,6 +305,8 @@ def my_main(epochs):
         loss = loss_fn(y_pred, y_train.squeeze(-1))
         print_(loss.item())
 
+        mlflow.log_param("loss", loss.item)
+
         # Zero gradients
         optimizer.zero_grad()
         loss.backward()  # Gradients
@@ -328,5 +331,4 @@ def my_main(epochs):
 
 
 with mlflow.start_run() as run:
-    my_main(epochs)
-    mlflow.log_param("epochs", epochs)
\ No newline at end of file
+    my_main(epochs)
\ No newline at end of file