From 127487369f534414778d98fdb19da20566fdb5f1 Mon Sep 17 00:00:00 2001 From: Zhun Zhong Date: Sun, 26 May 2019 15:37:19 +1000 Subject: [PATCH] Fix bug for prefetcher Set input and target to Cuda when without using prefetcher. --- train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/train.py b/train.py index 63003763..38d86992 100644 --- a/train.py +++ b/train.py @@ -428,6 +428,9 @@ def validate(model, loader, loss_fn, args): with torch.no_grad(): for batch_idx, (input, target) in enumerate(loader): last_batch = batch_idx == last_idx + if not args.prefetcher: + input = input.cuda() + target = target.cuda() output = model(input) if isinstance(output, (tuple, list)):