Depends on what you mean by swapping out. You should be able to do fast asynchronous host2device and device2host copies if you use a background stream (torch.cuda.Stream) and a pinned buffer for staging. But if I were you I’d make sure it’s really worth it. You probably don’t want to spend too much time debugging synchronization issues just to save a few ms.
Regular slicing is very fast. Returned tensor shares memory with the original one, only has some metadata modified.
The memory is ready to be reused as soon as the tensor goes out of scope in Python (unless you have reference cycles in your objects). The GPU memory is only growing because we’re caching CUDA allocations. cudaMalloc and cudaFree are slow and synchronize CPU with the device, so it’s unacceptable to call them all the time. Our CUDA allocator will keep all the allocated blocks around, because it’ll assume that you’re going to use them at the next iteration too. However, if you’re going to operate on inputs of different lengths, you might want to do a warm up pass with the largest length you’re going to use throught the training. This will ensure that the allocated buffers will be able to contain any of the sequences. Otherwise the allocator has to do some compacting and allocate larger blocks once you use a sequence that doesn’t fit into the blocks you have allocated.