fix14
This commit is contained in:
parent
3e53e29514
commit
d07246f63f
@ -28,5 +28,5 @@ RUN chmod +x ./train.py
|
|||||||
ARG epochs=5
|
ARG epochs=5
|
||||||
RUN echo $epochs
|
RUN echo $epochs
|
||||||
|
|
||||||
CMD python3 ./train.py epochs
|
CMD python ./train.py epochs
|
||||||
#CMD python3 ./test.py
|
#CMD python3 ./test.py
|
2
train.py
2
train.py
@ -74,7 +74,7 @@ if __name__ == '__main__':
|
|||||||
criterion = nn.CrossEntropyLoss()
|
criterion = nn.CrossEntropyLoss()
|
||||||
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
|
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
|
||||||
|
|
||||||
epochs = sys.argv[1]
|
epochs = int(sys.argv[1])
|
||||||
print(epochs)
|
print(epochs)
|
||||||
|
|
||||||
trainNet(trainloader, criterion, optimizer, int(float(epochs)))
|
trainNet(trainloader, criterion, optimizer, int(float(epochs)))
|
||||||
|
Loading…
Reference in New Issue
Block a user