diff --git a/Jenkinsfile_train b/Jenkinsfile_train
index 541bc47..8bb39d4 100644
--- a/Jenkinsfile_train
+++ b/Jenkinsfile_train
@@ -33,6 +33,7 @@ pipeline {
                     img.inside {
                         sh 'chmod +x dlgssdpytorch.py'
                         sh 'python3 ./dlgssdpytorch.py $PARAMETRY'
+						sh 'rm -rf my_model'
                         sh 'chmod +x train_mlflow.py'
 						//sh 'chmod +x generate_MLmodel.py'
 						//sh 'python3 ./generate_MLmodel.py'
@@ -46,7 +47,7 @@ pipeline {
 		stage('end') {
             steps {
 				//Zarchiwizuj wynik
-				archiveArtifacts 'model.pt, my_runs/**/*'
+				archiveArtifacts 'model.pt, my_runs/**/*, my_model/**/*'
 					
 			}
 		}
diff --git a/train_mlflow.py b/train_mlflow.py
index f293ca6..75fb2c0 100644
--- a/train_mlflow.py
+++ b/train_mlflow.py
@@ -10,7 +10,6 @@ import mlflow.pytorch
 from urllib.parse import urlparse
 from mlflow.models.signature import infer_signature
 
-
 class LayerLinearRegression(nn.Module):
     def __init__(self):
         super().__init__()
@@ -27,6 +26,7 @@ parser.add_argument('-e', '--epochs', type=int, default=100, help="Liczba epok",
 args = parser.parse_args()
 
 if __name__ == "__main__":
+
     lr = args.lr
     n_epochs = args.epochs
     mlflow.set_experiment("s426206")
@@ -133,7 +133,9 @@ if __name__ == "__main__":
         tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
         # print(tracking_url_type_store)
         # Model registry does not work with file store
+
         if tracking_url_type_store != "file":
-            mlflow.sklearn.log_model(model, "model", registered_model_name="s426206", signature=siganture, input_example=input_example)
+            mlflow.pytorch.log_model(model, "model", registered_model_name="s426206", signature=siganture, input_example=input_example)
         else:
-            mlflow.sklearn.log_model(model, "model", signature=siganture, input_example=input_example)
\ No newline at end of file
+            mlflow.pytorch.log_model(model, "model", signature=siganture, input_example=input_example)
+            mlflow.pytorch.save_model(model, "my_model", signature=siganture, input_example=input_example)
\ No newline at end of file