import torch


def predict(test_loader, device, model, model_type, dim):
    """
    Parameters
    ----------
    test_loader : torch.utils.data.DataLoader
        data to predict on
    device : str
        'cuda' or 'cpu'
    model : torch.nn.Module
        model containing the weights to compute the predictions
    model_type : str
        'regression' or 'classification', for the latter an arg max is applied
    dim : int
        number of time series

    Returns
    -------
    torch.tensor
    """
    pred = []
    for test_batch in test_loader:
        test_batch = test_batch.to(device)
        with torch.no_grad():
            out = model(test_batch[:, :, :dim*2], test_batch[:, :, -1])
            if model_type=='classification':
                out = out.argmax(1)
        pred.append(out)
    return torch.cat(pred, 0).squeeze()