This commit is contained in:
James Betker 2021-06-04 21:26:54 -06:00
parent 6c8c8087d5
commit 7d45132f60

View File

@ -52,7 +52,8 @@ class CategorizationLossEvaluator(evaluator.Evaluator):
self.model.eval()
with torch.no_grad():
for hq, labels in tqdm(self.dataloader):
for batch in tqdm(self.dataloader):
hq, labels = batch['hq'], batch['labels']
hq = hq.to(self.env['device'])
labels = labels.to(self.env['device'])
if self.masking: