from typing import Union
import torch


def tensor_container_to_device(obj: Union[torch.Tensor, tuple, list, dict], device: str, non_blocking=True):
    """
    recursively send compounded objects to device (sending all tensors to device and maintaining structure)
        :param obj           the object to send to device (list / tuple / tensor / dict)
        :param device:       device to send the tensors to
        :param non_blocking: used for DistributedDataParallel
        :returns        an object with the same structure (tensors, lists, tuples) with the device pointers (like
                        the return value of Tensor.to(device)
    """
    if isinstance(obj, torch.Tensor):
        return obj.to(device, non_blocking=non_blocking)
    elif isinstance(obj, tuple):
        return tuple(tensor_container_to_device(x, device, non_blocking=non_blocking) for x in obj)
    elif isinstance(obj, list):
        return [tensor_container_to_device(x, device, non_blocking=non_blocking) for x in obj]
    elif isinstance(obj, dict):
        return {k: tensor_container_to_device(v, device, non_blocking=non_blocking) for k, v in obj.items()}
    else:
        return obj