PyTorch首个Epoch训练后触发CUDA device-side assert错误求解

• 优先确认损失函数与模型输出层的匹配性：
  二分类场景下如果使用CrossEntropyLoss，模型最后一层输出维度必须为2，输入的标签需为torch.long类型的0/1值；如果使用BCEWithLogitsLoss，模型最后一层输出维度必须为1，输入的标签需转成torch.float32类型，且形状与输出张量完全对齐。第一个epoch未触发报错属于偶发情况，第二个epoch遍历到形状、类型不匹配的样本时就会触发底层断言。
• 排查自定义Dataset的返回值：
  检查__getitem__方法的返回逻辑，确保不存在偶发的标签值超出0/1范围、标签带有冗余维度的问题，每个样本返回的标签必须是合法的0/1标量。
• 检查设备一致性：
  确保每个batch的输入张量、标签、模型、损失函数都部署在同一个设备上，避免出现部分张量在CPU、部分在CUDA的情况。第一个epoch结束后的验证阶段如果将部分张量意外移到CPU，第二轮训练开始前没有移回CUDA也会触发该错误。
• 开启CUDA同步报错定位问题根因：
  在代码最开头加入以下配置，让报错指向真正触发异常的代码行，而非作为异步报错同步点的running_loss计算行：

  import os
  os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
  

• 验证训练流程的梯度逻辑：
  确保每个batch训练前都执行了optimizer.zero_grad()操作，梯度累积异常经过多轮迭代后也可能触发数值越界。

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