From 2350b7aca0f5c6fae11fc08d2624748c85571fea Mon Sep 17 00:00:00 2001 From: Xinqi Fan Date: Wed, 21 Apr 2021 12:35:09 +0800 Subject: [PATCH] Pass batch size and workers to val loader pass opt.batch_size and opt.workers to val_loader --- main_ce.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/main_ce.py b/main_ce.py index 29573d30..cbe88254 100644 --- a/main_ce.py +++ b/main_ce.py @@ -161,8 +161,8 @@ def set_loader(opt): train_dataset, batch_size=opt.batch_size, shuffle=(train_sampler is None), num_workers=opt.num_workers, pin_memory=True, sampler=train_sampler) val_loader = torch.utils.data.DataLoader( - val_dataset, batch_size=256, shuffle=False, - num_workers=8, pin_memory=True) + val_dataset, batch_size=opt.batch_size, shuffle=False, + num_workers=opt.num_workers, pin_memory=True) return train_loader, val_loader