I have found non_blocking=True to be very dangerous when going from GPU->CPU. For example:
import torch
action_gpu = torch.tensor([1.0], device=torch.device('cuda'), pin_memory=True)
print(action_gpu)
action_cpu = action_gpu.to(torch.device('cpu'), non_blocking=True)
print(action_cpu)

output
tensor([1.], device='cuda:0')
tensor([0.])

Process finished with exit code 0

Any idea why the tensors are not equal? I would expect the thread to block until the transfer from the GPU is finished.