pytorch model代码内tensor device不一致的问题

在编写一段处理两个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)