Skip to content
Snippets Groups Projects

change async to non_blocking

Merged Saeed SARFJOO requested to merge fix_py37_compatibility into master
1 file
+ 1
1
Compare changes
  • Side-by-side
  • Inline
@@ -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()
Loading