import torch


def get_xy_map(out_h, out_w):
    """
    created coordinate map with the given shape.
    Returns map of shape (out_h, out_w, 2)
    """
    y_map = torch.arange(out_h).float() / (out_h - 1)
    x_map = torch.arange(out_w).float() / (out_w - 1)
    x_map = x_map.expand(out_h, *x_map.shape)
    y_map = y_map.expand(out_w, *y_map.shape).transpose(1, 0)
    xy_map = torch.stack((x_map, y_map), dim=2)
    return xy_map