执行numba.cuda.close()后，如何让PyTorch重新调用GPU（无需重启内核）

我之前也踩过这个坑！当你在同一个Python内核里混合使用numba的CUDA和PyTorch时，两者会共享系统的CUDA上下文——直接调用cu.close()或者device.reset()会粗暴地销毁这个上下文，但PyTorch后续还会尝试使用已失效的上下文，就会抛出RuntimeError: CUDA error: invalid argument这类错误。不用重启内核的话，用PyTorch自带的上下文重置方法就能恢复，具体操作如下：

核心解决方案 #

通过PyTorch自身的设备重置接口，重新创建有效的CUDA上下文，覆盖被numba搞乱的环境：

torch.cuda.device(你的设备号，比如0).reset()

修正后的完整测试代码 #

from numba import cuda as cu
import torch

# 初始状态验证：PyTorch正常访问GPU
a = torch.rand(100, 100)
print("初始状态：a.cuda()成功，设备为", a.cuda().device)

# 模拟问题场景：调用numba的CUDA操作后尝试关闭
try:
    device = cu.get_current_device()
    cu.close()
except RuntimeError as e:
    print(f"调用cu.close()时触发预期错误：{e}")

# 此时PyTorch调用cuda会报错，执行恢复操作
torch.cuda.device(0).reset()

# 验证恢复效果
a = torch.rand(100, 100)
print("恢复后：a.cuda()成功，设备为", a.cuda().device)
print("torch.cuda.is_available()结果：", torch.cuda.is_available())

背后的原因 #

numba和PyTorch在同一个进程里会共享CUDA上下文资源：

• numba的cu.close()会直接销毁当前活跃的CUDA上下文，但没有通知PyTorch
• PyTorch后续的CUDA操作依然绑定在已销毁的上下文上，自然触发无效参数错误
• torch.cuda.device().reset()会让PyTorch主动创建一个全新的、有效的CUDA上下文，重新获得GPU的控制权

额外注意事项 #

• 尽量避免在同一个内核中混用numba CUDA和PyTorch CUDA，如果必须混用，建议用PyTorch的上下文管理器（比如with torch.cuda.device(0):）来确保上下文一致性
• 不要随意调用numba的device.reset()或cu.close()，除非你确定要彻底终止CUDA会话并准备好重新初始化

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