最近在调一个图像分割的模型，训练倒是正常，但到了验证阶段每次加载权重后一跑前向就OOM。我用的是PyTorch 2.0，显卡是4090 24G，batch size已经降到1了，输入图片也resize到256x256。诡异的是，训练时同样的batch size和分辨率都没事，就验证的时候爆。我怀疑是不是我用了torch.no_grad()但忘记把模型切到eval()模式，导致BN层还在更新？还是说DataLoader的pin_memory=True在验证时反而占用更多显存？另外，我加载的是保存的完整checkpoint（包含optimizer状态），会不会是optimizer的动量缓冲也占了一部分显存？有遇到过类似情况的朋友吗？求指点一下排查思路。