mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
Fix bug for prefetcher
Set input and target to Cuda when without using prefetcher.
This commit is contained in:
parent
e8cf619005
commit
127487369f
3
train.py
3
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)):
|
||||
|
Loading…
x
Reference in New Issue
Block a user