import torch


def flatten_feature_map(feature_map):
    """
    Flatten input tensor from [batch, y, x, f] to [batch, y*x*f]
    Args:
        feature_map: tensor with shape [batch, y, x, f]
    Returns:
        reshaped_map: tensor with  shape [batch, y*x*f]
    """
    map_shape = feature_map.get_shape()
    if(map_shape.ndims == 4):
        (batch, y, x, f) = map_shape
        prev_input_features = int(y * x * f)
        resh_map  = torch.reshape(feature_map, [-1, prev_input_features])
    elif(map_shape.ndims == 2):
        resh_map = feature_map
    else:
      raise ValueError("Input feature_map has incorrect ndims")
    return resh_map