PyTorch uses a memory cache to avoid malloc/free calls and tries to reuse the memory, if possible, as described in the docs.
To release memory from the cache so that other processes can use it, you could call torch.cuda.empty_cache().
EDIT: sorry, just realized that you are already using this approach. I’ll try to reproduce the observation.
EDIT2: the code seems to work fine and on my system only the CUDA context remains on the device.