include progress bar for test

NTU24pg · Nov 10, 2020 · f40ac62 · f40ac62
1 parent a3fa286
commit f40ac62
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/processor.py b/src/processor.py
@@ -4,6 +4,8 @@
 from .star import STAR
 from .utils import *
 
+from tqdm import tqdm
+
 
 class processor(object):
     def __init__(self, args):
@@ -154,7 +156,7 @@ def test_epoch(self):
         error_epoch, final_error_epoch = 0, 0,
         error_cnt_epoch, final_error_cnt_epoch = 1e-5, 1e-5
 
-        for batch in range(self.dataloader.testbatchnums):
+        for batch in tqdm(range(self.dataloader.testbatchnums)):
 
             inputs, batch_id = self.dataloader.get_test_batch(batch)
             inputs = tuple([torch.Tensor(i) for i in inputs])