From 5c792ac14ca67dfaa76198ce890a0f28b8ac4ca1 Mon Sep 17 00:00:00 2001
From: Adrian Charkiewicz <riraasaa@gmail.com>
Date: Sun, 8 May 2022 21:40:38 +0200
Subject: [PATCH] save only model, not everything

---
 Jenkinsfile3       | 2 +-
 pytorch/pytorch.py | 7 ++++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/Jenkinsfile3 b/Jenkinsfile3
index 0c00834..0b3bab9 100644
--- a/Jenkinsfile3
+++ b/Jenkinsfile3
@@ -24,7 +24,7 @@ stages {
                     archiveArtifacts artifacts: 'pytorch/result.txt',
                    allowEmptyArchive: true
 				   dir('my_runs') {
-                    archiveArtifacts artifacts: '**/**'
+                    archiveArtifacts artifacts: '**/*.pth'
                 }
                
             }
diff --git a/pytorch/pytorch.py b/pytorch/pytorch.py
index e3a4c6a..3a6ca07 100644
--- a/pytorch/pytorch.py
+++ b/pytorch/pytorch.py
@@ -170,7 +170,10 @@ def fit(epochs, lr, model, train_loader, val_loader, _log, opt_func=torch.optim.
         result = evaluate(model, val_loader)
         model.epoch_end(epoch, result, epochs)
         history.append(result)
-    ex.add_artifact("saved_model.pb")
+    
+    torch.save(model, 'saved_model.pth')
+    ex.add_artifact("saved_model.pth")
+    
     return history
 
 
@@ -207,5 +210,7 @@ def main():
     #my_config()
     #print("number of epochs is: ", epochs)
     history5 = fit()
+    
+   
 
 #ex.run()