Хороший гайд про то как устроены non_blocking и pin_memory операции в PyTorch, наглядно демонстрирует как работает host2device, device2host копирование и рассказывает в каких ситуациях стоит использовать non_blocking и pin_memory, а в каких нет. https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html