save only model, not everything

2022-05-08 21:40:38 +02:00 · 2022-05-08 21:40:38 +02:00 · 5c792ac14c
parent 9ce0127a72
commit 5c792ac14c
2 changed files with 7 additions and 2 deletions
--- a/2
+++ b/2
@ -24,7 +24,7 @@ stages {
                    archiveArtifacts artifacts: 'pytorch/result.txt',
                   allowEmptyArchive: true
 				   dir('my_runs') {
-                    archiveArtifacts artifacts: '**/**'
+                    archiveArtifacts artifacts: '**/*.pth'
                }
               
            }
--- 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


@ -208,4 +211,6 @@ def main():
    #print("number of epochs is: ", epochs)
    history5 = fit()
    
+   
+
 #ex.run()