diff --git a/Biblioteka_DL/dllib-mlflow.py b/Biblioteka_DL/dllib-mlflow.py
index 2be6024..a57aef9 100644
--- a/Biblioteka_DL/dllib-mlflow.py
+++ b/Biblioteka_DL/dllib-mlflow.py
@@ -296,34 +296,34 @@ def my_main(epochs):
         print ("The loss calculated: ", loss)
 
 
-    # Not using dataloader
-    x_train, y_train = Variable(torch.from_numpy(features_train_g)).float(), Variable(torch.from_numpy(labels_train_g)).long()
-    for epoch in range(1, epochs + 1):
-        print("Epoch #", epoch)
-        y_pred = model(x_train)
+    with mlflow.start_run() as run:
+    	x_train, y_train = Variable(torch.from_numpy(features_train_g)).float(), Variable(torch.from_numpy(labels_train_g)).long()
+    	for epoch in range(1, epochs + 1):
+        	print("Epoch #", epoch)
+        	y_pred = model(x_train)
 
-        loss = loss_fn(y_pred, y_train.squeeze(-1))
-        print_(loss.item())
+        	loss = loss_fn(y_pred, y_train.squeeze(-1))
+        	print_(loss.item())
 
-        # Zero gradients
-        optimizer.zero_grad()
-        loss.backward()  # Gradients
-        optimizer.step()  # Update
+        	# Zero gradients
+        	optimizer.zero_grad()
+        	loss.backward()  # Gradients
+        	optimizer.step()  # Update
 
-    # Prediction
-    x_test = Variable(torch.from_numpy(features_test_g)).float()
-    pred = model(x_test)
+    	# Prediction
+    	x_test = Variable(torch.from_numpy(features_test_g)).float()
+    	pred = model(x_test)
 
-    pred = pred.detach().numpy()
+    	pred = pred.detach().numpy()
 
-    print("The accuracy is", accuracy_score(labels_test_g, np.argmax(pred, axis=1)))
-    mlflow.log_metric("accuracy", accuracy_score(labels_test_g, np.argmax(pred, axis=1)))
+    	print("The accuracy is", accuracy_score(labels_test_g, np.argmax(pred, axis=1)))
+    	mlflow.log_metric("accuracy", accuracy_score(labels_test_g, np.argmax(pred, axis=1)))
 
-    pred = pd.DataFrame(pred)
+    	pred = pd.DataFrame(pred)
 
-    pred.to_csv('result.csv')
+    	pred.to_csv('result.csv')
 
-    # save model
-    torch.save(model, "games_model.pkl")
+    	# save model
+    	torch.save(model, "games_model.pkl")