from typing import Any
import torch


def move_to(obj: Any, device: torch.device):
    """Credit: https://discuss.pytorch.org/t/pytorch-tensor-to-device-for-a-list-of-dict/66283
    Arguments:
        obj {dict, list} -- Object to be moved to device
        device {torch.device} -- Device that object will be moved to
    Raises:
        TypeError: object is of type that is not implemented to process
    Returns:
        type(obj) -- same object but moved to specified device
    """
    if torch.is_tensor(obj) or isinstance(obj, torch.nn.Module):
        return obj.to(device)
    if isinstance(obj, dict):
        res = {k: move_to(v, device) for k, v in obj.items()}
        return res
    if isinstance(obj, list):
        return [move_to(v, device) for v in obj]
    if isinstance(obj, tuple):
        return tuple(move_to(list(obj), device))
    
    return obj