import torch


def run_inference_vanilla(tokens_tensor, segments_tensor, model):
  """Run inference on the model.

  Args:
    tokens_tensor: The tokens to infer the activation from.
    segments_tensor: Segments of the sequence to retrieve the activation for.
    model: The model to run inference on.

  Returns:
    layers_act: Activation of the model given the parameters.
  """
  with torch.no_grad():
    layers_act, _ = model(tokens_tensor, segments_tensor)
    # Reorder to get more intuitive ordering of batches, layers, words, neurons
    layers_act = torch.stack(layers_act).permute(1, 0, 2, 3)
  return layers_act