How to move multiple tensors to the Cuda device concurrently?
policy_data, value_data, action_mask = policy_data.cuda(non_blocking=True), value_data.cuda(non_blocking=True), action_mask.cuda(non_blocking=True)

rewards, regret_probs = rewards.cuda(non_blocking=True), regret_probs.cuda(non_blocking=True)

return action_probs.cpu(), sample_probs.cpu(), sample_indices.cpu(), update

I am doing some RL work and am wondering whether it would be possible to speed up fragments like the above by launching the data transfer to the GPU on different streams before waiting on them together. Does PyTorch have any functions that would make this easier? I'd rather ask here before I dive into the minutiae of optimizing data transfers.

Seems like one potential solution would be to pack all of the data into a single tensor (though of course you'd likely pay a small cost due to unused elements within this compacted representation.) An alternative would be to store this compact tensor as a sparse tensor (no additional data, but slightly more memory consumption per value). You'd have to test between these two to determine which was more efficient for your use case.
