fix14

2023-05-07 18:34:53 +02:00 · 2023-05-07 18:34:53 +02:00 · d07246f63f
commit d07246f63f
parent 3e53e29514
2 changed files with 2 additions and 2 deletions
--- a/2
+++ b/2
@ -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
--- 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)))