update dllib-mlflow.py

Biblioteka_DL/dllib-mlflow.py

@@ -305,13 +305,13 @@ 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
         optimizer.step()  # Update
 
+    mlflow.log_param("loss", loss.item)
+
     # Prediction
     x_test = Variable(torch.from_numpy(features_test_g)).float()
     pred = model(x_test)