def spatial_average(in_tens, keepdim=True):
    """Takes the average over non-batch and channel dimensions of the input.

    :param in_tens: 1 x 1 [x T] x H x W
    :param keepdim: Whether to retain the dimensions of the input
    """
    B, C = in_tens.size(0), in_tens.size(1)
    in_tens_reshaped = in_tens.view(B, C, -1)
    in_tens_mean = in_tens_reshaped.mean(dim=-1)
    if keepdim:
        in_tens_mean = in_tens_mean.view(*tuple([1 for _ in range(len(in_tens.size()))]))
    return in_tens_mean