def _data_to_batch(data: list[list[int]], max_seq_len: int,
                   pad_token_id: int) -> dict[str, torch.Tensor]:
    """Helper function to create a proper batch of data."""
    input_ids = torch.stack([
        torch.tensor(d + [pad_token_id] * (max_seq_len - len(d))) for d in data
    ])

    attention_mask = torch.stack([
        torch.tensor([1] * len(d) + [pad_token_id] * (max_seq_len - len(d)))
        for d in data
    ])
    return {'input_ids': input_ids, 'attention_mask': attention_mask}
Helper function to create a proper batch of data.
tests/data/test_packing.py