Show progress bar per epoch

nshaud · May 30, 2018 · 5eada10 · 5eada10
1 parent afb6081
commit 5eada10
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models.py b/models.py
@@ -1119,7 +1119,7 @@ def train(net, optimizer, criterion, data_loader, epoch, scheduler=None,
         avg_loss = 0.
 
         # Run the training loop for one epoch
-        for batch_idx, (data, target) in enumerate(data_loader):
+        for batch_idx, (data, target) in tqdm(enumerate(data_loader), total=len(data_loader)):
             # Load the data into the GPU if required
             if cuda:
                 data, target = data.cuda(), target.cuda()