From 09d0026265fda359e35db6fdec969e6fae534466 Mon Sep 17 00:00:00 2001
From: Maciej Czajka <maccza5@st.amu.edu.pl>
Date: Wed, 11 May 2022 15:21:21 +0200
Subject: [PATCH] update dllib-mlflow.py

---
 Biblioteka_DL/dllib-mlflow.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

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