torch.concat()结果存于CPU？GPU运算设备不一致原理问询

首先看你代码里的关键问题：你只把c和d移到了GPU，但s2 = torch.concat((a, b), dim=2)里用的是a和b——这两个变量你没在给出的代码里定义并迁移到GPU，大概率是之前在CPU上创建的张量，所以torch.concat的结果自然会和输入张量（a、b）保持同一设备，也就是CPU。

背后的原理： #

• PyTorch中，张量操作的结果设备默认与输入张量的设备一致。如果所有输入张量都在同一个设备上，输出也会在该设备；如果输入张量设备不同，默认会直接报错（不会自动跨设备迁移）。
• 你的s1 = c+d中，c和d都是GPU张量，所以计算结果s1也在GPU；而s2的输入是CPU上的a和b，结果自然在CPU。

修正方案： #

把代码里的a和b换成你已经移到GPU的c和d，或者提前将a、b也迁移到GPU：

c = torch.rand((2000, 64, 64)).to('cuda')
d = torch.rand((2000, 64, 64)).to('cuda')
t3 = time.time()
s1 = c+d
# 修正变量名，使用已迁移到GPU的c和d
s2 = torch.concat((c, d), dim=2)
t4 = time.time()

这样s2的设备就会和c、d一致，都是GPU了。

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