load image problem

[StannyHo]

I inference my data set, but when I want to load model , it cause error in this function

def bilinear_sampler(img, coords, mode='bilinear', mask=False):
    """ Wrapper for grid_sample, uses pixel coordinates """
    H, W = img.shape[-2:]
    xgrid, ygrid = coords.split([1,1], dim=-1)
    xgrid = 2*xgrid/(W-1) - 1
    ygrid = 2*ygrid/(H-1) - 1

    grid = torch.cat([xgrid, ygrid], dim=-1)
    img = F.grid_sample(img, grid, align_corners=True)

    if mask:
        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
        return img, mask.float()

    return img

Error : RuntimeError: grid_sampler(): expected grid and input to have same batch size, but got input with sizes [760, 1, 19, 40] and grid with sizes [720, 9, 9, 2]
How to fix it!!!
Thank you