import torch


def flip_feature_maps(feature_maps, flip_index=None):
    """Flip the feature maps and swap the channels.

     Args:
        feature_maps (list(torch.Tensor)): Feature maps.
        flip_index (list(int) | None): Channel-flip indexes. If None,
            do not flip channels.
    Returns:
        flipped_feature_maps (list(torch.Tensor)): Flipped feature_maps.
    """
    flipped_feature_maps = []
    for feature_map in feature_maps:
        feature_map = torch.flip(feature_map, [3])
        if flip_index is not None:
            flipped_feature_maps.append(feature_map[:, flip_index, :, :])
        else:
            flipped_feature_maps.append(feature_map)

    return flipped_feature_maps