import torch


def _get_grid_locations(image_height, image_width, device):
    """Wrapper for `torch.meshgrid`."""
    y_range = torch.arange(0., image_height, device=device, requires_grad=False)
    x_range = torch.arange(0., image_width, device=device, requires_grad=False)
    y_grid, x_grid = torch.meshgrid(y_range, x_range)
    return torch.stack((y_grid, x_grid), dim=-1)