import torch


def nested_to_device(data, device, non_blocking=True):
    """
        transfer data to device. The data is a nested structure
    Args:
        data (Dict, List, Tuple, torch.Tensor): nested data.
        device (torch.device): The target device.

    Kwargs: 
        non_blocking (Bool): Wether to block.

    Returns: The same structure of data. Each Tensor is transfered to `device`.
    """
    if isinstance(data, torch.Tensor):
        return data.to(device, non_blocking=non_blocking)
    elif isinstance(data, dict):
        return {k: nested_to_device(v, device, non_blocking) for k, v in data.items()}
    elif isinstance(data, (list, tuple)):
        return [nested_to_device(v, device, non_blocking) for v in data]
    else:
        raise ValueError(f'Type {type(data)} is not supported.')