From d07246f63fa57fdca896c45f8e8e33d2f945c6ac Mon Sep 17 00:00:00 2001
From: Jakub Henyk <jakubhenyk@gmail.com>
Date: Sun, 7 May 2023 18:34:53 +0200
Subject: [PATCH] fix14

---
 Dockerfile | 2 +-
 train.py   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index edd189c..2519e2d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -28,5 +28,5 @@ RUN chmod +x ./train.py
 ARG epochs=5
 RUN echo $epochs
 
-CMD python3 ./train.py epochs
+CMD python ./train.py epochs
 #CMD python3 ./test.py
\ No newline at end of file
diff --git a/train.py b/train.py
index 8700078..101d867 100644
--- a/train.py
+++ b/train.py
@@ -74,7 +74,7 @@ if __name__ == '__main__':
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
 
-    epochs = sys.argv[1]
+    epochs = int(sys.argv[1])
     print(epochs)
 
     trainNet(trainloader, criterion, optimizer, int(float(epochs)))