Merge pull request #2 from zhunzhong07/patch-1

Fix bug for prefetcher
2019-05-26 09:31:40 -07:00 · 2019-05-26 09:31:40 -07:00 · 337702d51b
parent e8cf619005 127487369f
commit 337702d51b
1 changed files with 3 additions and 0 deletions
--- 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)):