Skip to content
Snippets Groups Projects
Commit 57ebc7ce authored by Saeed SARFJOO's avatar Saeed SARFJOO
Browse files

change async to non_blocking

parent 5878e035
Branches
Tags
1 merge request!39change async to non_blocking
Pipeline #33388 passed
......@@ -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()
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment