from typing import Dict
from typing import List
import torch


def convert_alignments_to_tensor(
    alignments: Dict[str, List[int]], device: torch.device
) -> Dict[str, torch.Tensor]:
    """Convert alignments from list of int to a 1-D torch.Tensor.

    Args:
      alignments:
        A dict containing alignments. Keys are utterance IDs and
        values are their corresponding frame-wise alignments.
      device:
        The device to move the alignments to.
    Returns:
      Return a dict using 1-D torch.Tensor to store the alignments.
      The dtype of the tensor are `torch.int64`. We choose `torch.int64`
      because `torch.nn.functional.one_hot` requires that.
    """
    ans = {}
    for utt_id, ali in alignments.items():
        ali = torch.tensor(ali, dtype=torch.int64, device=device)
        ans[utt_id] = ali
    return ans