基于demo.py进行batch推理

[Redtides0]

你好作者，感谢你的工作，多视角视频生成的方向的工作经常采用gim来计算不同视角视频同一时间步下的对应帧之间的match pix num来衡量视角一致性，我基于demo.py做了一下修改打算实现batch推理，使用的模型是DKMV3，以下是修改的地方

# 多图读取
img_path0 = "assets/demo/a1.png"
img_path1 = "assets/demo/a2.png"
img_path2 = "assets/demo/b1.png"
img_path3 = "assets/demo/b2.png"

image0 = read_image(img_path0)
image1 = read_image(img_path1)
image0, scale0 = preprocess(image0)
image1, scale1 = preprocess(image1)

image2 = read_image(img_path2)
image3 = read_image(img_path3)
image2, scale2 = preprocess(image2)
image3, scale3 = preprocess(image3)

image0 = image0.to(device)[None]
image1 = image1.to(device)[None]
image2 = image2.to(device)[None]
image3 = image3.to(device)[None]
# 构建batch
image0 = torch.cat((image0, image2), dim=0)
image1 = torch.cat((image1, image3), dim=0)



with warnings.catch_warnings():
   warnings.simplefilter("ignore")
   # 推理时设置batch为True，upsample_preds为False
   model.upsample_preds = False
   dense_matches, dense_certainty = model.match(image0_, image1_, batched=True)
   sparse_matches, mconf = model.sample(dense_matches, dense_certainty, 5000)

但是代码会抛出以下错误
Traceback (most recent call last):
File "/home/zirui/projects/gim/demo.py", line 452, in
dense_matches, dense_certainty = model.match(image0_, image1_, batched=True)
File "/home/zirui/projects/gim/networks/dkm/models/dkm.py", line 707, in match
dense_certainty = dense_certainty - low_res_certainty
RuntimeError: The size of tensor a (1333) must match the size of tensor b (896) at non-singleton dimension 3

请问应该如何解决？