在编写一段处理两个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
这段代码报了这个错误文章来源:https://www.toymoban.com/news/detail-539386.html
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!
统一下进行掩码计算的张量的设备即可文章来源地址https://www.toymoban.com/news/detail-539386.html
device = mask.Device
indices = indices.to(device)
到了这里,关于pytorch model代码内tensor device不一致的问题的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!