I managed to find the culprit of the problem. It’s not actually related to any of the code I posted above, but something upstream. But it does appear that torch.cuda.empty_cache() cannot clean all cached memory. I’ve searched through most of the documentations available, and the best I got is

Calling  empty_cache()  22 releases all  unused  cached memory from PyTorch so that those can be used by other GPU applications

I’m not quite sure what kind of cached memory is used. It’ll be great if someone can clarify