def spatial_average(in_tens, keepdim=True):
    """Returns the mean value of each row of the input tensor in the spatial
    dimension.

    Args:
        in_tens (Tensor): Tensor with shape [N, C, H, W].
        keepdim (bool, optional): If keepdim is True, the output tensor is of
            the shape [N, C, 1, 1]. Otherwise, the output will have shape
            [N, C]. Defaults to True.

    Returns:
        Tensor: Tensor after average pooling to 1x1 with shape [N, C, 1, 1] or
            [N, C].
    """
    return in_tens.mean([2, 3], keepdim=keepdim)