import torch


def apply_mask(hidden_states, masks):
    """Apply mask to masked tokens in a batch

    Args:
        hidden_states (torch.tensor): shape(batch_size, seq_len(PADDED), hidden_dim)
        masks (torch.tensor): shape(batch_size, seq_len(PADDED)) where masked tokens=0, unmasked tokens=1

    Returns:
        masked_hidden_states (torch.tensor): shape(batch_size, seq_len(PADDED), hidden_dim)
    """
    hidden_dim = hidden_states.shape[-1]
    hidden_states.view(-1, hidden_dim)[~masks.view(-1).type(torch.ByteTensor), :] = 0
    return hidden_states