如何使用Torch清除指定CUDA设备的缓存（DDP场景）

要清除指定GPU的缓存，核心思路是先切换到目标CUDA设备的上下文，再调用torch.cuda.empty_cache()——这个方法本身就是针对当前活跃的CUDA设备生效的。具体有两种实用实现方式：

1. 使用上下文管理器（推荐） #

用torch.cuda.device()上下文管理器临时切换到目标设备，执行缓存清理后自动回到原设备，不会干扰后续代码的设备上下文：

# 清除GPU 3的缓存
with torch.cuda.device('cuda:3'):
    torch.cuda.empty_cache()

2. 手动切换设备 #

如果需要在一段代码内持续操作目标设备，可以先手动切换，清理完成后再切回原设备：

# 记录当前活跃设备
original_device = torch.cuda.current_device()
# 切换到GPU 3
torch.cuda.set_device(3)
# 清除缓存
torch.cuda.empty_cache()
# 切回原设备
torch.cuda.set_device(original_device)

针对DDP场景的额外提醒 #

结合你用DDP的场景，补充两点：

• 如果辅助模型运行在绑定GPU 3的独立进程中，直接在该进程内执行上述缓存清理代码即可，无需跨进程操作。
• 若在同一个进程内管理多GPU模型，务必确保清理缓存前，辅助模型相关的未使用张量已经被正确释放（比如用del删除引用、解除变量绑定），否则empty_cache()无法释放被张量占用的显存。

另外，GPU内存不足的根源可能是内存泄漏，建议排查：

• 辅助模型的推理/训练循环中是否存在重复创建张量、未及时释放中间变量的情况
• 是否有全局变量持有大量张量引用，导致无法被垃圾回收机制回收

内容的提问来源于stack exchange，提问作者DareGhost