在Colab Pro运行KiU-Net 3D训练时遇Adam CUDA图捕获错误

错误原因 #

这个错误是因为Adam优化器初始化时未开启capturable=True参数，但训练流程中触发了PyTorch的CUDA图捕获机制，两者不兼容导致的。

修复步骤 #

• 打开train.py文件，找到初始化Adam优化器的代码段（通常类似如下结构）：

  opt = torch.optim.Adam(model.parameters(), lr=0.001)  # 你的学习率参数可能不同
  

• 在优化器初始化参数中添加capturable=True，修改后代码如下：

  opt = torch.optim.Adam(model.parameters(), lr=0.001, capturable=True)
  

• 保存修改后的train.py，重新执行!python train.py启动训练。

补充说明 #

如果训练代码中搭配了学习率调度器（如lr_scheduler），无需额外调整调度器代码，开启capturable=True后即可与多数调度器兼容。

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