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)))