I want to hide the time needed to load tensors from CPU to GPU. This can in general be done by using gpu_tensor = tensor.cuda(non_blocking=True).
This will block as soon as I access the gpu_tensor. But let’s suppose I do some work only after I access the tensor the first time on the GPU. To still be able to hide the CPU-GPU-movement time, I would like to load the tensor one iteration ahead.
Here is a little code example:
    tensors = [torch.rand((20000, 800), dtype=torch.float32).pin_memory(),
               torch.rand((20000, 800), dtype=torch.float32).pin_memory()]

    gpu_tensors = [torch.rand((20000, 800), dtype=torch.float32, device="cuda"),
                   torch.rand((20000, 800), dtype=torch.float32, device="cuda")]

    if pre_load:
        gpu_tensors[0].copy_(tensors[0], non_blocking=False)
    for batch_index in range(num_batches):
        to_gpu_time -= time.time()
        if pre_load:
            gpu_tensors[(batch_index+1)%2].copy_(tensors[(batch_index+1)%2], non_blocking=True)
        else:
            gpu_tensor = tensors[0].to(device, non_blocking=True)
        to_gpu_time += time.time()

        tensor_access_time -= time.time()
        if pre_load:
            gpu_tensor = gpu_tensors[batch_index%2]
        sum_result = gpu_tensor.sum().item()
        tensor_access_time += time.time()

        # do some work
        total_sleep_2_time -= time.time()
        if sleep_2:
            time.sleep(sleep_time)
        total_sleep_2_time += time.time()

        if write_back:
            write_back_time -= time.time()
            cpu_tensor = gpu_tensor.cpu()
            write_back_time += time.time()

In theory it should be possible to hide the CPU-GPU-movement time in the sleep time after the tensor access if we are pre-loading one tensor ahead. But unfortunately this seems not to work this way.
Do you have an idea why I am still losing time in the tensor_access part if I pre-load a tensor or do you know how I could solve this problem?