Fix bug for prefetcher

Set input and target to Cuda when without using prefetcher.
pull/2/head
Zhun Zhong 6 years ago committed by GitHub
parent e8cf619005
commit 127487369f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -428,6 +428,9 @@ def validate(model, loader, loss_fn, args):
with torch.no_grad(): with torch.no_grad():
for batch_idx, (input, target) in enumerate(loader): for batch_idx, (input, target) in enumerate(loader):
last_batch = batch_idx == last_idx last_batch = batch_idx == last_idx
if not args.prefetcher:
input = input.cuda()
target = target.cuda()
output = model(input) output = model(input)
if isinstance(output, (tuple, list)): if isinstance(output, (tuple, list)):

Loading…
Cancel
Save