I have changed this line
scaling_mat = torch.eye(3, device=frame.device).repeat(batch_size, 1, 1)

to
scaling_mat = torch.eye(256, device=frame.device).repeat(batch_size, 1, 1)

but there is a error  as follow.
Traceback (most recent call last):
File “ioucompute.py”, line 197, in 
target_mask = warp.warp_image(fake_template, scaling_mat, out_shape=fake_template.shape[-2:])
File “warp_image.py”, line 46, in warp_image
xy_warped = torch.matmul(H, xy)  # H.bmm(xy)
RuntimeError: Expected tensor to have size 256 at dimension 1, but got size 3 for argument #2 ‘batch2’ (while checking arguments for bmm)
The “warp_image” function is as follow.
def warp_image(img, H, out_shape=None, input_grid=None):
if out_shape is None:
out_shape = img.shape[-2:]
if len(img.shape) < 4:
img = img[None]
if len(H.shape) < 3:
H = H[None]
assert img.shape[0] == H.shape[0], ‘batch size of images do not match the batch size of homographies’
batchsize = img.shape[0]
# create grid for interpolation (in frame coordinates)
if input_grid is None:
y, x = torch.meshgrid([
torch.linspace(-utils.BASE_RANGE, utils.BASE_RANGE,
steps=out_shape[-2]),
torch.linspace(-utils.BASE_RANGE, utils.BASE_RANGE,
steps=out_shape[-1])
])
x = x.to(img.device)
y = y.to(img.device)
else:
x, y = input_grid
x, y = x.flatten(), y.flatten()
# append ones for homogeneous coordinates
xy = torch.stack([x, y, torch.ones_like(x)])
xy = xy.repeat([batchsize, 1, 1])  # shape: (B, 3, N)
# warp points to model coordinates
xy_warped = torch.matmul(H, xy)  # H.bmm(xy)
xy_warped, z_warped = xy_warped.split(2, dim=1)

# we multiply by 2, since our homographies map to
# coordinates in the range [-0.5, 0.5]
xy_warped = 2.0 * xy_warped / (z_warped + 1e-8)
x_warped, y_warped = torch.unbind(xy_warped, dim=1)
# build grid
grid = torch.stack([
    x_warped.view(batchsize, *out_shape[-2:]),
    y_warped.view(batchsize, *out_shape[-2:])
],
    dim=-1)

warped_img = torch.nn.functional.grid_sample(
    img, grid, mode='bilinear', padding_mode='zeros')
batchsize0 = img.shape[0]
batchsize1 = img.shape[1]
batchsize2 = img.shape[2]
if utils.hasnan(warped_img):
    print('nan value in warped image! set to zeros')
    warped_img[utils.isnan(warped_img)] = 0
return warped_img