现象

在昇腾NPU卡上训练，不管是否在代码中添加with autocast()，显存占用都没变化。

原因分析

原因是，错误地使用了torch.cuda.amp.autocast，这个调用在cuda环境能生效，但在昇腾环境不生效，错误示例如下：

from torch.cuda.amp import autocastwith autocast(enabled=True):result = model(batch)...

解决办法是改成使用torch.autocast，示例如下：

with torch.autocast(enabled=True, device_type="npu", dtype=torch.float16):result = model(batch)

其它

笔者耗费了一整天才发现的解决方法，如果帮到你了请点赞。