import torch


def batch_mask_make(batch, device):
    """Create an attention mask
    Args:
        batch: (seq_len, batch_size)
        device: GPU or CPU device tensor
    
    Returns:
        attn_mask: (batch_size, seq_len)
    """
    attn_mask = torch.where(batch.permute(1, 0) != 0,
                            torch.ones((batch.size()[1], batch.size()[0]), device=device),
                            torch.zeros((batch.size()[1], batch.size()[0]), device=device))
    return attn_mask.to(device)