pytorch to_device遇到数据迁移不成功的问题

pytorch中使用to(device)后数据未迁移成功

问题描述

调试模型时遇到报错：

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking arugment for argument index in method wrapper_gather)

内容说是要进行运算的两个数据没有在同一个设备上，在cpu和cuda:0两个设备上。
首先可以推断是有个别数据未迁移到同一设备上，这里我选择cuda:0作为实验的设备，以加快训练速度。

错误代码

主调代码：

print(f"self.device:{self.device}")   # cuda:0
event_ix = torch.LongTensor(self.lm.get_events_in_sentence(relpair))
event_ix.to(self.device)   # 数据迁移到GPU
print(f"exp_HGRU event_ix.device:{event_ix.device}")
time_anchor = self.timeAnchorModel(s, s_a_mask, event_ix, self.device)

被调代码：

event_1_ix, event_2_ix = event_ix.split(1, dim=-1)
print(f"event_1_ix.device:{event_1_ix.device}\nevent_2_ix.device:{event_2_ix.device}") # cpu cpu
event_1 = torch.gather(sequence_output, dim=1,
                               index=event_1_ix.expand(batch_size, hidden_size).unsqueeze(dim=1))  # index处发生报错
event_2 = torch.gather(sequence_output, dim=1,
                               index=event_2_ix.expand(batch_size, hidden_size).unsqueeze(dim=1))

我把device通过参数传进来了，但是event_1_ix.device和event_2_ix.device仍然是cpu

解决方法

其实很简单，就是data.to(device)后需要赋值给另一变量，可以是data本身，否者无法真正实现迁移。
可能是因为简单，网上没有针对性的解决方法，对pytorch略懂皮毛的人来说，有些简单的点，却最磨人。

正确代码

event_ix = torch.LongTensor(self.lm.get_events_in_sentence(relpair)).to(self.device)

或者

event_ix = torch.LongTensor(self.lm.get_events_in_sentence(relpair))
event_ix = event_ix.to(self.device)

！data = data.to(device) ------ 需要data来承接