from torch.amp import autocast ModuleNotFoundError: No module named 'torch.amp'

### 解决 `ModuleNotFoundError: No module named 'torch.ao'` 和 `torch.amp.autocast` 的问题

遇到 `ModuleNotFoundError: No module named 'torch.ao'` 这类错误通常是因为使用的 PyTorch 版本不支持该模块或存在路径拼写错误。对于特定于 `torch.amp.autocast` 的情况，这可能意味着当前环境中安装的 PyTorch 版本较旧或是安装过程存在问题。

#### 验证并更新 PyTorch 安装

为了确保可以正常使用 `torch.amp.autocast` 功能，建议验证现有 PyTorch 安装是否是最新的稳定版。如果版本过低，则应考虑升级至最新版本[^1]：

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

此命令会根据官方推荐配置，在 conda 虚拟环境下重新安装兼容 CUDA 11.3 的 PyTorch 及其依赖项。

#### 正确引入 `autocast`

确认 PyTorch 已正确安装后，应该能够顺利导入 `torch.cuda.amp.autocast`。注意这里不是 `torch.ao`，而是 `torch.cuda.amp` 下的一个子功能[^4]：

from torch.cuda.amp import autocast

此外，当使用自动混合精度训练模型时，除了 `autocast` 外还经常配合 `GradScaler` 来实现更稳定的梯度缩放操作。

#### 使用示例

以下是利用 `autocast` 和 `GradScaler` 实现自动混合精度训练的小例子：

import torch
from torch.cuda.amp import GradScaler, autocast

model = ...  # 初始化模型
optimizer = ...  # 设置优化器
scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = criterion(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

通过上述调整，应当能有效解决因 PyTorch 版本不符而导致无法找到相应模块的问题，并成功应用自动混合精度加速 GPU 上的计算任务。