在编写一段处理两个tensor的代码如下,需要在forward函数内编写函数创建一个新的tensor进行索引的掩码计算
# todo(liang)空间交换
def compute_sim_and_swap(t1, t2, threshold=0.7):
n, c, h, w = t1.shape
sim = torch.nn.functional.cosine_similarity(t1, t2, dim=1) # n, h, w
sim = sim.unsqueeze(0) # c, n, h, w
expand_tensor = sim.clone()
# 使用拼接构建相同的维度
for _ in range(c-1): # c, n, h, w
sim = torch.cat([sim, expand_tensor], dim=0)
sim = sim.permute(1, 0, 2, 3) # n, c, h, w
# 创建逻辑掩码,小于 threshold 的将掩码变为 True 用于交换
mask = sim < threshold
indices = torch.rand(mask.shape) < 0.5
t1[mask&indices], t2[mask&indices] = t2[mask&indices], t1[mask&indices]
return t1, t2
这段代码报了这个错误
File "xxx/network.py", line 347, in compute_sim_and_swap
t1[mask&indices], t2[mask&indices] = t2[mask&indices], t1[mask&indices]
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
统一下进行掩码计算的张量的设备即可
device = mask.Device
indices = indices.to(device)