def forward(self, *input, **kwargs):
        """Condition the Flamingo layers on the media locations before forward
        function."""
        input_ids = kwargs['input_ids'] if 'input_ids' in kwargs else input[0]
        media_locations = input_ids == self.media_token_id
        if self.only_attend_previous:
            attend_previous = True
        elif self.use_media_placement_augmentation:
            attend_previous = (random.random() < 0.5)
        else:
            attend_previous = False

        for layer in self.get_decoder().layers:
            layer.condition_media_locations(media_locations)
            layer.condition_attend_previous(attend_previous)

        return super().forward(
            *input, **kwargs)  # Call the other parent's forward method
Condition the Flamingo layers on the media locations before forward
        function.
mmpretrain/models/multimodal/flamingo/adapter.py