from typing import Dict
import torch


def to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
    """
    Sends a batch of data to the given torch devicee (cpu or cuda).

    Args:
        batch (Dict[str, torch.Tensor]): Batch to be send to the device.
        device (torch.device): Device (either torch.device('cpu') or torch.device('cuda').

    Returns:
        Dict[str, torch.Tensor]: The batch at the given device.

    """

    return {key: val.to(device) for key, val in batch.items()}