updated train.py

IshanRattan · Aug 7, 2024 · a8bcd54 · a8bcd54
1 parent ce921e8
commit a8bcd54
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/dl/image/classification/cifar10/train.py b/dl/image/classification/cifar10/train.py
@@ -31,8 +31,12 @@
 # # model(x)
 # # print(model)
 criterion = nn.CrossEntropyLoss()
-optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate, momentum=config.momentum)
-_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=config.scheduler_step_size, gamma=config.scheduler_gamma)
+optimizer = torch.optim.SGD(model.parameters(),
+                            lr=config.learning_rate,
+                            momentum=config.momentum)
+_lr_scheduler = lr_scheduler.StepLR(optimizer,
+                                    step_size=config.scheduler_step_size,
+                                    gamma=config.scheduler_gamma)
 model, losses = train_model(model,
                        criterion,
                        optimizer,