import torch


def _forward_step(model, batch, device):
    """Run a forward step of model for a given batch of data

    Parameters
    ----------
    model : torch.nn.Module
      torch model, an instance of torch.nn.Module
    batch : list
      a 2 element list of inputs and labels, to be fed to the model

    Returns
    -------
    torch.tensor
      output of model's forward function
    """

    # put model in eval mode
    model.eval()
    model.to(device)

    with torch.no_grad():
        # inputs and targets
        inputs = batch[0]
        # move data to DEVICE
        if len(inputs) == 1:
            inputs = inputs.to(device)
            r = model(inputs)
        else:
            for item in inputs:
                if isinstance(item, torch.Tensor):
                    item = item.to(device)
            r = model(*inputs)
        # forward
        return r