import torch


def predict_fn(input_data, model):
    """ Predict stock prices
        Parameters:
        ----------
        input_data(num of samples, lookback): input data
        model: trained model
        
        Returns:
        --------
        output (num of samples, output_size): predictions
    """
    
    print('Predicting stock prices')

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    data = torch.from_numpy(input_data).float()
    data = data.to(device)

    # Make sure to put the model into evaluation mode
    model.eval()

    output = model.forward(data)
    
    out_np = output.cpu().detach().numpy()
    
    return out_np