pytorch 1.7.0 torchvision 0.8.1 torch.cuda.amp gradscaler DDP 训练卡死

报错：pytorch DDP 模型卡住
代码

# 具体卡住的代码
yolov5训练代码 train.py 中有一句：

                scaler.step(optimizer)  # optimizer.step

程序运行到第二个epoch的时候，卡住了，
具体卡在调用语句：
/home/xxx/lib/python3.7/site-packages/torch/cuda/amp/grad_scaler.py

中的
        if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):

就卡死了。。。。并不只知道为什么

软件环境
python 3.7.11
torch 1.7.0
torchvision 0.8.1
cuda 10.1.243
nvidia驱动 450.80.02
pip 22.0.4

硬件环境
teslaP40 4GPU
2080ti 也会卡住

解决：

# 升级软件

$ pip install torch==1.8.0 # 更新torch
# pip卸载现有，
$ pip uninstall torchvision # 卸载现有torchvision
# 然后pip install torchvision， torch 直接升级到最新版本
$ pip install torchvision # 更新torchvision,会自动更新torch

----------- 最终版本------

安装torch
'1.11.0+cu102'
安装torchvision
'0.12.0+cu102'


注意：直接按照上述版本pip install 安装可能装不上。。