import torch


def grid_like(tensor):
  """
  Constructs a regular grid with the same spatial dimensions as the input tensor,
  on the same device. For an input of shape NxCxD_1x...xD_n returns grid of shape
  nxD_1x...xD_n, i.e. the size of the first dimension is equal to the number of
  spatial dimensions.

  :param tensor: tensor of reference shape/device
  :return: A regular grid
  """
  return torch.stack(torch.meshgrid([torch.arange(0, d, device=tensor.device) for d in tensor.shape[2:]]))