add parameter

Jenkinsfile

@@ -1,11 +1,13 @@
 pipeline {
     agent any
+    }
     parameters {
             buildSelector (
                 defaultSelector: lastSuccessful(),
                 description: 'Build for copying artifacts',
                 name: 'BUILD_SELECTOR'
             )
+            string(name: 'EPOCHS', defaultValue: '10', description: 'epochs')
         }
     stages {
         stage('Git Checkout') {
@@ -23,11 +25,10 @@ pipeline {
               script {
                def customImage = docker.build("custom-image")
                     customImage.inside {
-                        sh 'python3 ./model.py'
+                        sh 'python3 ./model.py ${params.EPOCHS}'
                         archiveArtifacts artifacts: 'model.pth, predictions.txt', onlyIfSuccessful: true
                     }
               }
             }
         }
-     }
-}
+     }

model.py

@@ -6,6 +6,7 @@ import pandas as pd
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import LabelEncoder
 import torch.nn.functional as F
+import os
  
 
 device = (
@@ -59,7 +60,7 @@ def main():
     val_loader = DataLoader(list(zip(X_val, y_val)), batch_size=64)
 
     # Training loop
-    epochs = 10
+    epochs = os.getenv("EPOCHS")
     for epoch in range(epochs):
         model.train()  # Set model to training mode
         running_loss = 0.0