PyTorch训练时X.to(device)无效引发输入与权重张量类型不匹配报错

PyTorch中张量的.to()方法不是原地修改操作，它会返回一个迁移到目标设备的新张量，原张量不会发生任何改变。你当前代码里只执行了X.to(device)、y.to(device)，没有接收返回的新张量，所以原X和y还是留在CPU上，自然X.is_cuda返回false，后续运行时就会出现输入和模型权重设备不匹配的报错。

只需要修改数据迁移的两行代码，接收.to()方法返回的新张量即可：

def train_loop(self, dataloader, device):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        print(device)
        # 接收迁移到GPU后的新张量
        X = X.to(device)
        print(X.is_cuda)
        y = y.to(device)
        
        pred = self.model(X)
        loss = self.loss_fn(pred, y)

• 不要混用.to(device)和硬编码的.cuda()写法，.to(device)兼容性更好，切换CPU/GPU训练不需要修改代码
• 如果使用DataLoader加载数据，可以在创建DataLoader时设置pin_memory=True，配合X = X.to(device, non_blocking=True)可以进一步提升数据迁移的效率
• 模型调用.to(device)是原地修改操作，不需要重新赋值，和张量的行为不同，注意区分避免混淆

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