Commit 57ebc7ce authored by Saeed SARFJOO's avatar Saeed SARFJOO

change async to non_blocking

parent 5878e035
Pipeline #33388 passed with stage
in 30 minutes and 18 seconds
......@@ -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