Commit fdabcc91 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

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): ...@@ -148,7 +148,7 @@ class DltResNetExtractor(Preprocessor):
with torch.no_grad(): with torch.no_grad():
in_batch = numpy.array(batch) in_batch = numpy.array(batch)
if cuda_flag == 1: 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: else:
input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1)) input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1))
fX = model(input_var).data.cpu().numpy() 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