From 829bdd40a6802e1e8c5717900b4bb6458610e702 Mon Sep 17 00:00:00 2001 From: Ming Du Date: Thu, 2 May 2024 12:12:56 -0500 Subject: [PATCH] Add tqdm to tester --- generic_trainer/tester.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/generic_trainer/tester.py b/generic_trainer/tester.py index 526e89f..5ee16d5 100644 --- 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)