Add tqdm to tester

AdvancedPhotonSource · May 2, 2024 · 829bdd4 · 829bdd4
1 parent 75be2ae
commit 829bdd4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/generic_trainer/tester.py b/generic_trainer/tester.py
@@ -4,6 +4,7 @@
 import numpy as np
 import torch
 from torch.utils.data import Dataset, DataLoader
+import tqdm
 
 import generic_trainer.trainer as trainer
 from generic_trainer.configs import *
@@ -94,7 +95,7 @@ def build_dir(self):
     def run(self):
         if self.mode == 'state_dict':
             self.model.eval()
-        for j, data_and_labels in enumerate(self.dataloader):
+        for j, data_and_labels in enumerate(tqdm.tqdm(self.dataloader)):
             data, labels = self.process_data_loader_yield(data_and_labels)
             if self.mode == 'state_dict':
                 preds = self.model(*data)