From 42ec57f385cc33876f940499afb974f1d5eef476 Mon Sep 17 00:00:00 2001 From: Jakub Henyk Date: Sun, 7 May 2023 18:26:05 +0200 Subject: [PATCH] fix9 --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index ec38ef2..d438e51 100644 --- a/train.py +++ b/train.py @@ -74,9 +74,9 @@ if __name__ == '__main__': criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) - epochs = int(sys.argv[1]) + epochs = sys.argv[1] - trainNet(trainloader, criterion, optimizer, epochs) + trainNet(trainloader, criterion, optimizer, int(epochs)) PATH = './cifar_net.pth' torch.save(net.state_dict(), PATH)