import torch


def predict_dataloader(model, dataloader, discard_target=True):
    """
    Return predictions for the given dataloader and model.
    
    Parameters
    ----------
    model : pytorch model
        The pytorch model to predict with.
    dataloader : pytorch dataloader
        Dataloader returning batches of inputs. Shuffle should be False if
        order is important in the predictions
    discard_target : bool (default = True)
        If True, onyl the first element of the batch is kept. This is useful if
        the dataloader returns batches as (inputs, targets, ...). If it only
        returns inputs directly as a tensor, set this to False.
    
    Returns
    -------
    predictions : list of tensors
        List of predicted batch.
    """
    model_device = next(model.parameters()).device
    predictions = []
    model.eval()
    with torch.no_grad():
        for batch in dataloader:
            if discard_target:
                batch = batch[0]
                
            batch = batch.to(model_device)
            predictions.append(model(batch))
    return predictions