import torch


def batch_step_mask(masked, decoded):
    """Create a new mask by adding a decoded position to an old mask.
        Arguments:
            masked: The old mask of shape (batch_size, max_len).
            decoded: The decoded positions of shape (batch_size,).
        Returns:
            A new mask of shape (batch_size, max_len).
    """
    decoded = decoded.unsqueeze(-1)
    return masked.scatter(-1, decoded, torch.zeros_like(decoded, dtype=masked.dtype))