PyTorch韩语聊天机器人GPU训练报错：设备不匹配问题求助

你的问题核心是训练时输入张量仍在CPU，而模型已移至GPU，导致设备不匹配。按以下步骤修改代码即可解决：

1. 将训练循环中的输入张量移至GPU #

在训练循环里，拿到token_ids, mask, label后，立即添加设备转移代码：

token_ids, mask, label = samples
# 新增：将所有输入张量移至指定设备
token_ids = token_ids.to(device)
mask = mask.to(device)
label = label.to(device)

2. 修正sneg的设备问题 #

sneg是Python浮点数，直接与GPU张量运算时可能隐式生成CPU张量，需将其转为GPU张量：

# 原代码
sneg = -1e18  
# 修改为
sneg = torch.tensor(-1e18, device=device)

3. 完整修改后的训练循环片段 #

print("::start::")
for epoch in tqdm(range(epoch)):
    for batch_idx, samples in enumerate(train_dataLoader):
        optimizer.zero_grad()
        token_ids, mask, label = samples
        # 新增设备转移
        token_ids = token_ids.to(device)
        mask = mask.to(device)
        label = label.to(device)
        
        out = model(token_ids)
        out = out.logits  
        mask_3d = mask.unsqueeze(dim=2).repeat_interleave(repeats=out.shape[2], dim=2)
        # 使用已移至GPU的sneg
        mask_out = torch.where(mask_3d == 1, out, sneg * torch.ones_like(out))
        loss = criterion(mask_out.transpose(2, 1), label)
        avg_loss = loss.sum() / mask.sum()
        avg_loss.backward()
        optimizer.step()
print("end")

额外检查点 #

• 确保aboutDataSets.ChatDataset返回的张量没有提前固定在CPU（本地能运行说明这块无问题，主要是Colab下需显式转移输入）
• 确认device已正确识别GPU：可在代码开头添加print(device)验证是否输出cuda:0

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