import torch


def prep_relevant_embeds(embeds: torch.Tensor, attention_mask: torch.Tensor):
    """Prepares the token embeddings for mean and max types

    Args:
        embeds: The token embeddings.
        attention_mask: The attention mask.

    Returns:
        torch.Tensor: The tensor with padding embeddings set to zero.

    """
    return embeds * attention_mask.unsqueeze(2).repeat(1, 1, embeds.shape[2])