Resolving epochs being hard-coded

microsoft · Oct 5, 2023 · 9052147 · 9052147
1 parent bfad08f
commit 9052147
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/training/cifar/cifar10_deepspeed.py b/training/cifar/cifar10_deepspeed.py
@@ -343,7 +343,7 @@ def create_moe_param_groups(model):
 # We simply have to loop over our data iterator, and feed the inputs to the
 # network and optimize.
 
-for epoch in range(2):  # loop over the dataset multiple times
+for epoch in range(args.epochs):  # loop over the dataset multiple times
 
     running_loss = 0.0
     for i, data in enumerate(trainloader):