import torch


def prepare_batch(sample, use_cuda=False):
    """Prepares a batch for training/validation by moving it to GPU

    Args:
        sample (dict | :obj:`torch.Tensor` | list): The batch to prepare
        use_cuda (bool): Move to cuda

    Returns:
        (dict | :obj:`torch.Tensor` | list): The batch, on GPU or CPU
    """

    def _move_to_cuda(maybe_tensor):
        if torch.is_tensor(maybe_tensor):
            return maybe_tensor.cuda()
        elif isinstance(maybe_tensor, dict):
            return {key: _move_to_cuda(value) for key, value in maybe_tensor.items()}
        elif isinstance(maybe_tensor, list):
            return [_move_to_cuda(x) for x in maybe_tensor]
        else:
            return maybe_tensor

    if use_cuda:
        return _move_to_cuda(sample)
    else:
        return sample