import torch


def batch_replace_token(batch, token_index, unk_index, device):
    """Replace an token index with an unk index
    Args:
        batch: (seq_len, batch_size)
        token_index (int): token index that will be replaced
        unk_index (int): unk token index
        device: GPU or CPU device tensor
    
    Returns:
        replaced_batch: ()
    """ 
    replaced_batch = torch.where(batch == token_index,
                                 torch.tensor([unk_index], device=device),
                                 batch)
    return replaced_batch