import torch


def safe_to_device(X, device):
    """Move the tensor or model to the device if it doesn't already live there.

    This function will avoid copying if the object already lives in the
    correct context. Works on PyTorch tensors and PyTorch models that inherit
    from the torch.nn.Module class.

    Parameters
    ----------
    X: torch.tensor
        The tensor to move to a device.

    device: str
        The PyTorch context to move the tensor to.

    Returns
    -------
    X: torch.tensor
        The tensor on the desired device.
    """

    if isinstance(X, torch.Tensor):
        if X.device == device:
            return X
        return X.to(device)

    elif isinstance(X, torch.nn.Module):
        if next(X.parameters()).device == device:
            return X
        return X.to(device)