Commit fdabcc91 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH

Merge branch 'fix_py37_compatibility' into 'master'

change async to non_blocking

Closes #24

See merge request !39
parents 5878e035 57ebc7ce
Pipeline #34467 passed with stages
in 17 minutes and 1 second
......@@ -148,7 +148,7 @@ class DltResNetExtractor(Preprocessor):
with torch.no_grad():
in_batch = numpy.array(batch)
if cuda_flag == 1:
input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1).cuda(async=True))
input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1).cuda(non_blocking=True))
else:
input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1))
fX = model(input_var).data.cpu().numpy()
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment